Update app.py

app.py

@@ -6,7 +6,7 @@ import numpy as np
 import pandas as pd
 import os
 import logging
-from datasets import Dataset
+from datasets import Dataset, DatasetDict
 
 # Configure logging
 logging.basicConfig(level=logging.DEBUG)
@@ -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"].astype(int).tolist()  # Convert labels to integers
+            "label": feedback_data["correct_label"].astype(int).tolist()  # Ensure labels are integers
         }
 
         dataset = Dataset.from_dict(dataset_dict)