Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr
+iimport gradio as gr
 import torch
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification, Trainer, TrainingArguments
 from PIL import Image
@@ -86,7 +86,7 @@ def retrain_model(feedback_data):
 
         dataset_dict = {
             "image": [load_image(f) for f in feedback_data["image_path"]],
-            "label": feedback_data["correct_label"].tolist()
+            "label": [str(lbl) for lbl in feedback_data["correct_label"].tolist()]
         }
 
         dataset = Dataset.from_dict(dataset_dict)
@@ -95,7 +95,7 @@ def retrain_model(feedback_data):
         # Preprocess the dataset
         def preprocess(examples):
             inputs = feature_extractor(images=examples["image"], return_tensors="pt")
-            inputs["labels"] = examples["label"]
+            inputs["labels"] = [int(lbl) for lbl in examples["label"]]
             return inputs
 
         dataset = dataset.with_transform(preprocess)