Update app.py

app.py

@@ -1,411 +1,27 @@
-import os, sys, glob
-
-import numpy as np
-import pandas as pd
-
-from tqdm.notebook import tqdm
-
-import PIL.Image
-
-import transformers
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification, TrainingArguments, Trainer
-
-import datasets
-from datasets import load_dataset, Features, Array3D, DatasetDict, ClassLabel
-
-import torch
-
-from torchvision.transforms import (
-    CenterCrop,
-    Compose,
-    Normalize,
-    RandomHorizontalFlip,
-    RandomResizedCrop,
-    Resize,
-    ToTensor,
-)
-
-import evaluate
-
+# Load model directly
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from PIL import Image # Import the Image module
+import torch # Import the torch module
 import streamlit as st
 
-#!git lfs install
-
-token = st.sidebar.text_input("Enter your Hugging Face token:", type="password")
-
-logged = False
-path_set = False
-
-if st.sidebar.button("Login"):
-        with open("token.txt", "w") as f:
-            f.write(token)
-
-        # Intentar iniciar sesión con huggingface-cli
-        with st.spinner("Logging in..."):
-          exit_code = os.system(f"huggingface-cli login --token {token}")
-
-          if exit_code != 0:
-              st.sidebar.error("Login failed. Please check your token and try again.")
-          else:
-              st.sidebar.success("Logged in successfully!")
-              logged = True
-
-labels = ["CM05",
-          "FACTURA",
-          "advertisement",
-          "handwritten",
-          "scientific_report",
-          "budget",
-          "scientific_publication",
-          "presentation",
-          "file_folder",
-          "memo",
-          "resume",
-          "invoice",
-          "letter",
-          "questionnaire",
-          "form",
-          "news_article"]
-
-NUM_OF_LABELS = len(labels)
-
-label2id, id2label = dict(), dict()
-
-for i, label in enumerate(labels):
-    label2id[label] = i
-    id2label[i] = label
-
-st.title("Document AI")
-
-parent_dir = st.text_input("Enter the directory path:")
-#parent_dir = "/content/docs"
-#parent_dir = r"C:\Users\Windows\Documents\AI Ollama\docs\docs"
-
-subfolders = ['CM05', 'FACTURA']
-
-selected_subfolder = st.sidebar.selectbox("Selecciona la subcarpeta", subfolders)
-
-all_files_loaded = False
-
-if parent_dir:
-    if not os.path.exists(parent_dir):
-        st.error(f"The directory {parent_dir} does not exist.")
-    else:
-        path_set = True
-        st.success("Directory path set successfully.")
-        uploaded_files = st.sidebar.file_uploader("Subir archivos", type=['jpg'], accept_multiple_files=True)
-
-        if path_set:
-          if st.sidebar.button("Cargar"):
-            if uploaded_files:
-              for file in uploaded_files:
-                  # Obtener el nombre del archivo y la extensión
-                  filename = file.name
-                  file_extension = filename.split(".")[-1].lower()
-
-                  subfolder_path = os.path.join(parent_dir, selected_subfolder)
-
-                  os.makedirs(subfolder_path, exist_ok=True)
-
-                  existing_files = os.listdir(subfolder_path)
-                  file_count = len(existing_files)
-
-                  new_filename = f"{selected_subfolder}_{file_count + 1}.{file_extension}"
-
-                  file_path = os.path.join(subfolder_path, new_filename)
-
-                  # unique_filename = f"uploaded_file_{hash(file.getvalue())}.{file_extension}"
-
-                  # unique_filename = f"new_filename.{file_extension}"
-
-                  #file_path = os.path.join(parent_dir, unique_filename)
-                  #file_path = os.path.join(parent_dir, selected_subfolder, unique_filename)
-                  file_path = os.path.join(subfolder_path, new_filename)
-                  with open(file_path, "wb") as f:
-                      f.write(file.getvalue())
-              st.success("Files uploaded successfully.")
-
-        if st.sidebar.button("Mostrar contenido del directorio"):
-          dir_contents = {}
-          for subfolder in subfolders:
-              subfolder_path = os.path.join(parent_dir, subfolder)
-              if os.path.exists(subfolder_path):
-                  dir_contents[subfolder] = os.listdir(subfolder_path)
-              else:
-                  dir_contents[subfolder] = []
-
-          st.sidebar.write("Contenido actual del directorio:")
-          st.sidebar.json(dir_contents)
-
-        if st.sidebar.button("Cargo Todos Los Archivos"):
-          all_files_loaded = True
-
-if all_files_loaded:
-          all_files = glob.glob(os.path.join(parent_dir, "**"), recursive=True)
-
-          dir_path = os.path.join(parent_dir, "*", "*.jpg")
-
-          files_and_name = glob.glob(dir_path)
-
-          st.write(f"Files found: {files_and_name}")
-
-          metadata = pd.DataFrame(files_and_name, columns=["file_path"])
-
-          metadata['file_name'] = metadata['file_path'].apply(lambda x: x.split("/")[-2] + "/" + x.split("/")[-1])
-
-          metadata['label'] = metadata['file_path'].apply(lambda x: x.split("/")[-2])
-          metadata['label'].replace(label2id, inplace=True)
-
-          metadata = metadata.drop(columns=["file_path"])
-
-          metadata_file_location = os.path.join(parent_dir, "metadata.csv")
-          metadata.to_csv(metadata_file_location, index=False)
-
-          st.write(metadata.tail())
-
-          dataset = load_dataset(parent_dir)
-
-          dataset = dataset.cast_column("label", ClassLabel(names=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]))
-
-          dataset
-
-          dataset['train'][1]
-
-          train_split = dataset['train'].train_test_split(train_size=0.80)
-
-          ds = DatasetDict({
-              'train' : train_split['train'],
-              'eval' : train_split['test']
-          })
-
-          st.title("Document AI Model Configuration")
-
-          #MODEL_VERSION = st.text_input("Enter the model version:")
-
-          MODEL_CKPT = "microsoft/dit-base"
-          MODEL_NAME = "Classifier_CM05-v2"
-          #MODEL_NAME = MODEL_CKPT.split("/")[-1] + "-Classifier_CM05" + "-" + MODEL_VERSION
-
-          # if MODEL_VERSION:
-          #   st.write(f"Model version: {MODEL_VERSION}")
-
-          MODEL_CKPT
-
-          NUM_OF_EPOCHS=18
-          LEARNING_RATE=5e-5
-
-          BATCH_SIZE=32
-          DEVICE = torch.device("cuda")
-
-          REPORTS_TO='tensorboard'
-          STRATEGY = "epoch"
-
-          feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_CKPT)
-          feature_extractor
-
-          # normalize
-          normalize = Normalize(mean=feature_extractor.image_mean, std=feature_extractor.image_std)
-
-          # train_transforms
-          train_transforms = Compose(
-                  [
-                      RandomResizedCrop((feature_extractor.size['height'], feature_extractor.size['width'])),# it was a list that used [], not ()
-                      RandomHorizontalFlip(),
-                      ToTensor(),
-                      normalize
-                  ]
-              )
-
-          # eval_transforms
-          val_transforms = Compose(
-                  [
-                      Resize((feature_extractor.size['height'], feature_extractor.size['width'])),
-                      CenterCrop((feature_extractor.size['height'], feature_extractor.size['width'])),
-                      ToTensor(),
-                      normalize,
-                  ]
-              )
-
-          def preprocess_train(example_batch):
-              """
-              Apply train_transforms across a batch
-              """
-              example_batch["pixel_values"] = [
-                  train_transforms(image.convert("RGB")) for image in example_batch["image"]
-              ]
-              return example_batch
-
-          def preprocess_val(example_batch):
-              """
-              Apply val_transforms across a batch
-              """
-              example_batch["pixel_values"] = [val_transforms(image.convert("RGB")) for image in example_batch["image"]]
-              return example_batch
-
-          ds['train'].set_transform(preprocess_train)
-          ds['eval'].set_transform(preprocess_val)
-
-          ds['train'][0]
-
-          model = AutoModelForImageClassification.from_pretrained(
-              MODEL_CKPT,
-              label2id=label2id,
-              id2label=id2label,
-              ignore_mismatched_sizes=True,
-          ).to(DEVICE)
-
-          from transformers import TrainingArguments
-
-          args = TrainingArguments(
-              MODEL_NAME,
-              remove_unused_columns=False,
-              evaluation_strategy=STRATEGY,
-              save_strategy=STRATEGY,
-              logging_strategy="steps",
-              logging_steps=8,
-              logging_first_step=True,
-              learning_rate=LEARNING_RATE,
-              per_device_train_batch_size=BATCH_SIZE,
-              per_device_eval_batch_size=BATCH_SIZE,
-              gradient_accumulation_steps=4,
-              num_train_epochs=NUM_OF_EPOCHS,
-              warmup_ratio=0.10,
-              report_to=REPORTS_TO,
-              hub_private_repo=True,
-              push_to_hub=True
-          )
-
-          def compute_metrics(p):
-              accuracy_metric = evaluate.load("accuracy")
-              accuracy = accuracy_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                references=p.label_ids)['accuracy']
-
-              ### ------------------- F1 scores -------------------
-
-              f1_score_metric = evaluate.load("f1")
-              weighted_f1_score = f1_score_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                          references=p.label_ids,
-                                                          average='weighted')["f1"]
-
-              micro_f1_score = f1_score_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                      references=p.label_ids,
-                                                      average='micro')['f1']
-
-              macro_f1_score = f1_score_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                      references=p.label_ids,
-                                                      average='macro')["f1"]
-
-              ### ------------------- recall -------------------
-
-              recall_metric = evaluate.load("recall")
-              weighted_recall = recall_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                      references=p.label_ids,
-                                                      average='weighted')["recall"]
-
-              micro_recall = recall_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                  references=p.label_ids,
-                                                  average='micro')["recall"]
-
-              macro_recall = recall_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                  references=p.label_ids,
-                                                  average='macro')["recall"]
-
-              ### ------------------- precision -------------------
-
-              precision_metric = evaluate.load("precision")
-              weighted_precision = precision_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                            references=p.label_ids,
-                                                            average='weighted')["precision"]
-
-              micro_precision = precision_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                        references=p.label_ids,
-                                                        average='micro')["precision"]
-
-              macro_precision = precision_metric.compute(predictions=np.argmax(p.predictions, axis=1),
-                                                        references=p.label_ids,
-                                                        average='macro')["precision"]
-
-              return {"accuracy" : accuracy,
-                      "Weighted F1" : weighted_f1_score,
-                      "Micro F1" : micro_f1_score,
-                      "Macro F1" : macro_f1_score,
-                      "Weighted Recall" : weighted_recall,
-                      "Micro Recall" : micro_recall,
-                      "Macro Recall" : macro_recall,
-                      "Weighted Precision" : weighted_precision,
-                      "Micro Precision" : micro_precision,
-                      "Macro Precision" : macro_precision
-                      }
-
-          def collate_fn(examples):
-              pixel_values = torch.stack([example['pixel_values'] for example in examples])
-              labels = torch.tensor([example["label"] for example in examples])
-              return {"pixel_values": pixel_values, "labels": labels}
-
-
-
-          st.write("Model Configuration")
-          st.write(model)
+st.title("Image Classification")
 
-          st.write("Training Arguments")
-          st.write(args)
+uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "png"])
 
-          st.write("Dataset")
-          st.write(ds)
+#image_path = "/content/cm5_2.jpg" # Store the path as a string
 
-          state = st.session_state.get('state', {'training': False, 'pushing': False})
 
-          try:
-              if st.sidebar.button("Train"):
-                  state['training'] = True  # Activar el estado de entrenamiento
-                  with st.spinner("Training..."):
-                      st.write("Training...")
-                      trainer = Trainer(
-                          model=model,
-                          args=args,
-                          train_dataset=ds['train'],
-                          eval_dataset=ds['eval'],
-                          tokenizer=feature_extractor,
-                          compute_metrics=compute_metrics,
-                          data_collator=collate_fn
-                      )
-                      train_results = trainer.train()
-                  st.write(train_results)
-                  st.success("Training completed!")
-                  state['training'] = False  # Desactivar el estado de entrenamiento
+processor = AutoImageProcessor.from_pretrained("mateoluksenberg/dit-base-Classifier_CM05")
+model = AutoModelForImageClassification.from_pretrained("mateoluksenberg/dit-base-Classifier_CM05")
 
-                  # Save model, metrics, and state after successful training
-                  trainer.save_model()
-                  trainer.log_metrics("train", train_results.metrics)
-                  trainer.save_metrics("train", train_results.metrics)
-                  trainer.save_state()
+image = Image.open(uploaded_file) # Load the image from the file path
 
-              if st.sidebar.button("Push"):
-                  state['pushing'] = True  # Activar el estado de envío
-                  with st.spinner("Pushing to Hub..."):
-                      try:
-                          # Ensure trainer is initialized if not done earlier
-                          if 'trainer' not in locals():
-                              trainer = Trainer(
-                                  model=model,
-                                  args=args,
-                                  train_dataset=ds['train'],
-                                  eval_dataset=ds['eval'],
-                                  tokenizer=feature_extractor,
-                                  compute_metrics=compute_metrics,
-                                  data_collator=collate_fn
-                              )
+inputs = processor(image, return_tensors="pt") # Pass the image object to the processor
 
-                          trainer.push_to_hub()
-                          st.success("Model pushed to the hub successfully!")
-                          state['pushing'] = False  # Desactivar el estado de envío
-                      except Exception as push_error:
-                          st.error(f"Error pushing model to hub: {str(push_error)}")
+with torch.no_grad(): # Use torch.no_grad() to disable gradient calculations
+    logits = model(**inputs).logits
 
-          except Exception as e:
-              st.error(f"Error during training: {str(e)}")
+# model predicts one of the 1000 ImageNet classes
+predicted_label = logits.argmax(-1).item()
 
-          if state['training']:
-              st.spinner("Training...").spinner_container.markdown('')
-          if state['pushing']:
-              st.spinner("Pushing to Hub...").spinner_container.markdown('')
+print(model.config.id2label[predicted_label])