Update Training script

train2.py

@@ -27,6 +27,9 @@ os.makedirs(OUTPUT_DIR, exist_ok=True)
 print("🔹 Loading dataset from 'data/' folder...")
 dataset = load_dataset("imagefolder", data_dir="data")
 print(f"🔹 Dataset loaded. Columns: {dataset['train'].column_names}")
+print(f"🔹 Dataset splits: {list(dataset.keys())}")
+print(f"🔹 Number of training samples: {len(dataset['train'])}")
+print(f"🔹 Number of validation samples: {len(dataset['validation'])}")
 
 # ----------------------------
 # PREPROCESSOR
@@ -35,18 +38,33 @@ print(f"🔹 Loading processor from {MODEL_NAME}...")
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
 
 def transform(example):
-    # Determine correct image column
+    # Determine image column
     image_column = "image" if "image" in example else list(example.keys())[0]
-    img = example[image_column]
-    if isinstance(img, str):  # if path, open it
-        img = Image.open(img).convert("RGB")
-    elif isinstance(img, Image.Image):
-        img = img.convert("RGB")
-    else:
-        raise ValueError(f"Unknown type for image: {type(img)}")
-    
-    inputs = processor(img, return_tensors="pt")
-    inputs["label"] = example["label"]
+    images = example[image_column]
+
+    # Ensure list
+    if not isinstance(images, list):
+        images = [images]
+
+    processed_images = []
+    for img in images:
+        if isinstance(img, str):
+            img = Image.open(img).convert("RGB")
+        elif isinstance(img, Image.Image):
+            img = img.convert("RGB")
+        else:
+            raise ValueError(f"Unknown type for image: {type(img)}")
+        processed_images.append(img)
+
+    # Convert images to tensors
+    inputs = processor(images=processed_images, return_tensors="pt")
+
+    # Handle labels
+    labels = example["label"]
+    if not isinstance(labels, list):
+        labels = [labels]
+    inputs["labels"] = torch.tensor(labels)
+
     return inputs
 
 print("🔹 Applying transform to dataset...")
@@ -80,7 +98,7 @@ def compute_metrics(eval_pred):
     return metrics
 
 # ----------------------------
-# TRAINING ARGS
+# TRAINING ARGUMENTS
 # ----------------------------
 args = TrainingArguments(
     output_dir=OUTPUT_DIR,
@@ -117,7 +135,7 @@ trainer.train()
 print("🔹 Training complete.")
 
 # ----------------------------
-# SAVE FINAL MODEL + LABELS
+# SAVE MODEL + PROCESSOR + LABELS
 # ----------------------------
 print("🔹 Saving final model and processor...")
 trainer.save_model(OUTPUT_DIR)