Upload 5 files

Adding training files

Dockerfile

@@ -0,0 +1,33 @@
+FROM python:3.10-slim
+
+# Create non-root user
+RUN adduser --disabled-password --gecos '' user
+USER user
+
+# Environment variables
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH \
+    PORT=7860
+
+WORKDIR $HOME/app
+
+# Copy requirements first (better for Docker layer caching)
+COPY --chown=user requirements.txt ./ 
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application
+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
+RUN chmod +x startup.sh
+
+# Start API
+CMD ["bash", "startup.sh"]

app.py

@@ -0,0 +1,48 @@
+# app.py
+import os, json
+from fastapi import FastAPI, UploadFile
+from transformers import AutoImageProcessor, BeitForImageClassification
+from PIL import Image
+import torch
+
+MODEL_DIR = "outputs/beit-retina"
+CLASSES = ["AMD","DMO","DR","GLC","HR","Normal"]
+
+app = FastAPI(title="Retina Disease Classifier")
+
+# Lazy load model & processor
+processor = None
+model = None
+
+def load_model():
+    global processor, model, CLASSES
+    processor = AutoImageProcessor.from_pretrained(MODEL_DIR)
+    model = BeitForImageClassification.from_pretrained(MODEL_DIR)
+    with open(os.path.join(MODEL_DIR, "labels.json")) as f:
+        CLASSES = json.load(f)
+
+@app.on_event("startup")
+def startup_event():
+    if os.path.exists(MODEL_DIR):
+        load_model()
+
+@app.post("/predict")
+async def predict(file: UploadFile):
+    if model is None:
+        return {"error": "Model not trained yet"}
+    img = Image.open(file.file).convert("RGB")
+    inputs = processor(images=img, return_tensors="pt")
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    probs = torch.softmax(logits, dim=1)[0].tolist()
+    pred_id = int(torch.argmax(logits, dim=1).item())
+    return {
+        "class_id": CLASSES[pred_id],
+        "probabilities": [{CLASSES[i]: float(p) for i, p in enumerate(probs)}]
+    }
+
+@app.post("/train")
+async def train_endpoint():
+    os.system("python train.py")  # blocking training run
+    load_model()
+    return {"status": "Training complete and model reloaded"}

requirements.txt

@@ -0,0 +1,12 @@
+torch==2.2.0+cpu
+torchvision==0.17.0+cpu
+transformers
+datasets
+accelerate
+scikit-learn
+fastapi
+uvicorn[standard]
+pillow
+pydantic==2.8.2
+python-multipart==0.0.9
+huggingface_hub==0.24.6

startup.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -euo pipefail
+# In HF Spaces with Docker, CUDA is available if a GPU is provisioned.
+exec uvicorn app:app --host 0.0.0.0 --port ${PORT:-7860}

train.py

@@ -0,0 +1,71 @@
+# train.py
+import os, json
+from transformers import AutoImageProcessor, BeitForImageClassification, TrainingArguments, Trainer
+from datasets import load_dataset
+from sklearn.metrics import accuracy_score, f1_score
+import numpy as np
+
+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")))
+
+
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    preds = np.argmax(logits, axis=1)
+    return {
+        "accuracy": accuracy_score(labels, preds),
+        "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))
+
+    model = BeitForImageClassification.from_pretrained(
+        MODEL_NAME,
+        num_labels=len(CLASSES),
+        id2label={i: c for i, c in enumerate(CLASSES)},
+        label2id={c: i for i, c in enumerate(CLASSES)}
+    )
+
+    args = TrainingArguments(
+        output_dir=output_dir,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=batch_size,
+        num_train_epochs=epochs,
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        metric_for_best_model="f1_weighted",
+        logging_steps=50,
+        report_to="none"
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["validation"],
+        tokenizer=processor,
+        compute_metrics=compute_metrics
+    )
+
+    trainer.train()
+    model.save_pretrained(output_dir)
+    processor.save_pretrained(output_dir)
+
+    with open(os.path.join(output_dir, "labels.json"), "w") as f:
+        json.dump(CLASSES, f)
+    print("✅ Training complete. Model saved at:", output_dir)
+
+if __name__ == "__main__":
+    train()