Training the Beit Model

.gitattributes

@@ -0,0 +1,4 @@
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,30 @@
+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=/root/.cache/huggingface \
+    TRANSFORMERS_CACHE=/root/.cache/huggingface/transformers \
+    TORCH_HOME=/root/.cache/torch
+
+# 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"}