Update Training script

train2.py

@@ -9,6 +9,7 @@ from transformers import (
     Trainer
 )
 from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
+from PIL import Image
 
 # ----------------------------
 # CONFIG
@@ -17,38 +18,51 @@ MODEL_NAME = "microsoft/beit-base-patch16-224"
 OUTPUT_DIR = os.environ.get("OUTPUT_DIR", os.path.expanduser("~/outputs/beit-retina"))
 NUM_CLASSES = 6   # retina disease classes
 
-# Make sure output directory exists
+print(f"🔹 OUTPUT_DIR set to: {OUTPUT_DIR}")
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 # ----------------------------
 # LOAD DATASET
 # ----------------------------
-# Example: Replace this with your retina dataset
-# You can load a Hugging Face dataset or your own image folder dataset
-# Dataset format: train/valid/test folders each containing subfolders by class name
+print("🔹 Loading dataset from 'data/' folder...")
 dataset = load_dataset("imagefolder", data_dir="data")
+print(f"🔹 Dataset loaded. Columns: {dataset['train'].column_names}")
 
 # ----------------------------
 # PREPROCESSOR
 # ----------------------------
+print(f"🔹 Loading processor from {MODEL_NAME}...")
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
 
 def transform(example):
-    inputs = processor(example["image"], return_tensors="pt")
+    # Determine correct 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"]
     return inputs
 
-# Map preprocessing
+print("🔹 Applying transform to dataset...")
 dataset = dataset.with_transform(transform)
+print("🔹 Transform applied successfully.")
 
 # ----------------------------
 # MODEL
 # ----------------------------
+print(f"🔹 Loading BEiT model ({MODEL_NAME}) with {NUM_CLASSES} classes...")
 model = BeitForImageClassification.from_pretrained(
     MODEL_NAME,
     num_labels=NUM_CLASSES,
     ignore_mismatched_sizes=True
 )
+print("🔹 Model loaded successfully.")
 
 # ----------------------------
 # METRICS
@@ -56,12 +70,14 @@ model = BeitForImageClassification.from_pretrained(
 def compute_metrics(eval_pred):
     logits, labels = eval_pred
     preds = logits.argmax(axis=-1)
-    return {
+    metrics = {
         "accuracy": accuracy_score(labels, preds),
         "precision": precision_score(labels, preds, average="macro"),
         "recall": recall_score(labels, preds, average="macro"),
         "f1": f1_score(labels, preds, average="macro"),
     }
+    print(f"🔹 Metrics computed: {metrics}")
+    return metrics
 
 # ----------------------------
 # TRAINING ARGS
@@ -78,6 +94,7 @@ args = TrainingArguments(
     logging_dir=os.path.join(OUTPUT_DIR, "logs"),
     push_to_hub=False
 )
+print("🔹 TrainingArguments configured.")
 
 # ----------------------------
 # TRAINER
@@ -90,19 +107,22 @@ trainer = Trainer(
     tokenizer=processor,
     compute_metrics=compute_metrics
 )
+print("🔹 Trainer created. Ready to train.")
 
 # ----------------------------
 # TRAIN
 # ----------------------------
+print("🔹 Starting training...")
 trainer.train()
+print("🔹 Training complete.")
 
 # ----------------------------
 # SAVE FINAL MODEL + LABELS
 # ----------------------------
+print("🔹 Saving final model and processor...")
 trainer.save_model(OUTPUT_DIR)
 processor.save_pretrained(OUTPUT_DIR)
 
-# Save class labels mapping
 labels = dataset["train"].features["label"].names
 with open(os.path.join(OUTPUT_DIR, "labels.json"), "w") as f:
     json.dump(labels, f)