Update Training script

Dockerfile

@@ -1,5 +1,11 @@
 FROM python:3.10-slim
 
+# Install system dependencies (for Pillow and general Python packages)
+USER root
+RUN apt-get update && \
+    apt-get install -y libjpeg-dev zlib1g-dev git && \
+    rm -rf /var/lib/apt/lists/*
+
 # Create non-root user
 RUN adduser --disabled-password --gecos '' user
 USER user
@@ -7,27 +13,29 @@ USER user
 # Environment variables
 ENV HOME=/home/user \
     PATH=/home/user/.local/bin:$PATH \
-    PORT=7860
+    PORT=7860 \
+    HF_HOME=/home/user/.cache/huggingface \
+    TRANSFORMERS_CACHE=/home/user/.cache/huggingface/transformers \
+    TORCH_HOME=/home/user/.cache/torch \
+    OUTPUT_DIR=/home/user/outputs/beit-retina
 
 WORKDIR $HOME/app
 
-# Copy requirements first (better for Docker layer caching)
+# Create necessary directories
+RUN mkdir -p $HF_HOME $TRANSFORMERS_CACHE $TORCH_HOME $OUTPUT_DIR
+
+# Copy requirements first for caching
 COPY --chown=user requirements.txt ./ 
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy the rest of the application
+# Copy app
 COPY --chown=user . .
 
-# Expose FastAPI default port for Hugging Face Spaces
-EXPOSE 7860
-
-# HF auth picked automatically from env (Spaces provides HF_TOKEN)
-ENV HF_HOME=/home/user/.cache/huggingface \
-    TRANSFORMERS_CACHE=/home/user/.cache/huggingface/transformers \
-    TORCH_HOME=/home/user/.cache/torch
-
-RUN mkdir -p $HF_HOME $TRANSFORMERS_CACHE $TORCH_HOME
+# Make startup script executable
 RUN chmod +x startup.sh
 
+# Expose port
+EXPOSE 7860
+
 # Start API
-CMD ["bash", "startup.sh"]
+CMD ["bash", "startup.sh"]

train.py

@@ -1,17 +1,27 @@
-# train.py
-import os, json
+import os
+import json
+from PIL import Image
 from transformers import AutoImageProcessor, BeitForImageClassification, TrainingArguments, Trainer
 from datasets import load_dataset
 from sklearn.metrics import accuracy_score, f1_score
 import numpy as np
 
+# -------------------------------
+# Config
+# -------------------------------
 CLASSES = ["AMD","DMO","DR","GLC","HR","Normal"]
 MODEL_NAME = "microsoft/beit-base-patch16-224"
 
-print("HOME dir:", os.environ.get("HOME"))
-print("HF cache:", os.environ.get("HF_HOME", os.path.join(os.environ["HOME"], ".cache", "huggingface")))
+# Output directory (from env or default)
+OUTPUT_DIR = os.environ.get("OUTPUT_DIR", os.path.join(os.environ["HOME"], "outputs/beit-retina"))
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+print("HOME dir:", os.environ.get("HOME"))
+print("HF cache:", os.environ.get("HF_HOME"))
 
+# -------------------------------
+# Metrics
+# -------------------------------
 def compute_metrics(eval_pred):
     logits, labels = eval_pred
     preds = np.argmax(logits, axis=1)
@@ -20,16 +30,21 @@ def compute_metrics(eval_pred):
         "f1_weighted": f1_score(labels, preds, average="weighted")
     }
 
-def train(output_dir="/outputs/beit-retina", train_dir="data/train", val_dir="data/val", epochs=5, batch_size=16):
-    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
-    dataset = load_dataset("imagefolder", data_dir={"train": train_dir, "validation": val_dir})
-
-    def transform(examples):
-        images = [processor(Image.open(p).convert("RGB"), return_tensors="pt")["pixel_values"][0] for p in examples["image"]]
-        return {"pixel_values": images}
-
-    dataset = dataset.cast_column("label", dataset["train"].features["label"].cast(type="ClassLabel", names=CLASSES))
+# -------------------------------
+# Preprocessing function
+# -------------------------------
+def transform(examples, processor):
+    """Converts image paths to pixel_values tensors."""
+    images = [processor(Image.open(p).convert("RGB"), return_tensors="pt")["pixel_values"][0] 
+              for p in examples["image"]]
+    return {"pixel_values": images}
 
+# -------------------------------
+# Training function
+# -------------------------------
+def train(train_dir="data/train", val_dir="data/val", epochs=5, batch_size=16):
+    # Load processor and model
+    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
     model = BeitForImageClassification.from_pretrained(
         MODEL_NAME,
         num_labels=len(CLASSES),
@@ -37,8 +52,18 @@ def train(output_dir="/outputs/beit-retina", train_dir="data/train", val_dir="da
         label2id={c: i for i, c in enumerate(CLASSES)}
     )
 
+    # Load dataset
+    dataset = load_dataset("imagefolder", data_dir={"train": train_dir, "validation": val_dir})
+
+    # Map transform over dataset
+    dataset = dataset.map(lambda x: transform(x, processor), batched=True)
+
+    # Ensure dataset returns PyTorch tensors
+    dataset.set_format(type="torch", columns=["pixel_values", "label"])
+
+    # Training arguments
     args = TrainingArguments(
-        output_dir=output_dir,
+        output_dir=OUTPUT_DIR,
         per_device_train_batch_size=batch_size,
         per_device_eval_batch_size=batch_size,
         num_train_epochs=epochs,
@@ -50,6 +75,7 @@ def train(output_dir="/outputs/beit-retina", train_dir="data/train", val_dir="da
         report_to="none"
     )
 
+    # Trainer
     trainer = Trainer(
         model=model,
         args=args,
@@ -59,13 +85,19 @@ def train(output_dir="/outputs/beit-retina", train_dir="data/train", val_dir="da
         compute_metrics=compute_metrics
     )
 
+    # Train
     trainer.train()
-    model.save_pretrained(output_dir)
-    processor.save_pretrained(output_dir)
 
-    with open(os.path.join(output_dir, "labels.json"), "w") as f:
+    # Save model and processor
+    model.save_pretrained(OUTPUT_DIR)
+    processor.save_pretrained(OUTPUT_DIR)
+
+    # Save labels
+    with open(os.path.join(OUTPUT_DIR, "labels.json"), "w") as f:
         json.dump(CLASSES, f)
-    print("✅ Training complete. Model saved at:", output_dir)
+
+    print("✅ Training complete. Model saved at:", OUTPUT_DIR)
+
 
 if __name__ == "__main__":
     train()