Fine-tuning the Beit pre-trained model

.gitattributes

@@ -1,3 +1,9 @@
+app.py filter=lfs diff=lfs merge=lfs -text
+data filter=lfs diff=lfs merge=lfs -text
+Dockerfile filter=lfs diff=lfs merge=lfs -text
+requirements.txt filter=lfs diff=lfs merge=lfs -text
+startup.sh filter=lfs diff=lfs merge=lfs -text
+train.py filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
 *.jpeg filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -1,30 +1,3 @@
-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"]
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c22cda261116e37bfe0d54cd627c4bae393711ecf9241ae0bcb0708505de385
+size 771

app.py

@@ -1,48 +1,3 @@
-# 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"}
+version https://git-lfs.github.com/spec/v1
+oid sha256:4af62f87d886a945ea83769a8503175aa3223acce7bc67bf04dc77b7466fec56
+size 1522

requirements.txt

@@ -1,12 +1,3 @@
-torch
-torchvision
-transformers
-datasets
-accelerate
-scikit-learn
-fastapi
-uvicorn[standard]
-pillow
-pydantic==2.8.2
-python-multipart==0.0.9
-huggingface_hub==0.24.6
+version https://git-lfs.github.com/spec/v1
+oid sha256:64cfb5f9134f154569f48ce0aa7eafa9030ca81176aa3fc6194d385e36cd8bfb
+size 173

startup.sh

@@ -1,4 +1,3 @@
-#!/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}
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2ddf251bb3ea2548afe28f1275b333f1c01dcc6ed0d8f248faa7f2b24c1dda9
+size 165

train.py

@@ -1,71 +1,3 @@
-# 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()
+version https://git-lfs.github.com/spec/v1
+oid sha256:09b73a425fc928a3220b6e598bd618b454ca8b8d34cba277608dcd01baf3ea73
+size 2529