Update Training script

app.py

@@ -1,6 +1,6 @@
 # app.py
-import os, json
-from fastapi import FastAPI, UploadFile
+import os, json, subprocess
+from fastapi import BackgroundTasks, FastAPI, UploadFile
 from transformers import AutoImageProcessor, BeitForImageClassification
 from PIL import Image
 import torch
@@ -28,6 +28,21 @@ def load_model():
         processor, model = None, None
         print(f"⚠️ Skipping model load: {e}")
 
+def run_training():
+    try:
+        result = subprocess.run(
+            ["python", "train2.py"],
+            capture_output=True,
+            text=True
+        )
+        if result.returncode == 0 and os.path.exists(MODEL_DIR):
+            load_model()
+            print("✅ Training complete and model reloaded")
+        else:
+            print("❌ Training failed:", result.stderr)
+    except Exception as e:
+        print("⚠️ Training exception:", str(e))
+
 @app.on_event("startup")
 def startup_event():
     if os.path.exists(MODEL_DIR):
@@ -54,7 +69,7 @@ async def predict(file: UploadFile):
     }
 
 @app.post("/train")
-async def train_endpoint():
-    os.system("python train.py")  # blocking training run
-    load_model()
-    return {"status": "Training complete and model reloaded"}
+async def train_endpoint(background_tasks: BackgroundTasks):
+    # Schedule the training in the background
+    background_tasks.add_task(run_training)
+    return {"status": "Training started in background"}

train2.py

@@ -0,0 +1,110 @@
+import os
+import json
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoImageProcessor,
+    BeitForImageClassification,
+    TrainingArguments,
+    Trainer
+)
+from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
+
+# ----------------------------
+# CONFIG
+# ----------------------------
+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
+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
+dataset = load_dataset("imagefolder", data_dir="data")
+
+# ----------------------------
+# PREPROCESSOR
+# ----------------------------
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+
+def transform(example):
+    inputs = processor(example["image"], return_tensors="pt")
+    inputs["label"] = example["label"]
+    return inputs
+
+# Map preprocessing
+dataset = dataset.with_transform(transform)
+
+# ----------------------------
+# MODEL
+# ----------------------------
+model = BeitForImageClassification.from_pretrained(
+    MODEL_NAME,
+    num_labels=NUM_CLASSES,
+    ignore_mismatched_sizes=True
+)
+
+# ----------------------------
+# METRICS
+# ----------------------------
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    preds = logits.argmax(axis=-1)
+    return {
+        "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"),
+    }
+
+# ----------------------------
+# TRAINING ARGS
+# ----------------------------
+args = TrainingArguments(
+    output_dir=OUTPUT_DIR,
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    learning_rate=5e-5,
+    per_device_train_batch_size=16,
+    per_device_eval_batch_size=16,
+    num_train_epochs=5,
+    weight_decay=0.01,
+    logging_dir=os.path.join(OUTPUT_DIR, "logs"),
+    push_to_hub=False
+)
+
+# ----------------------------
+# TRAINER
+# ----------------------------
+trainer = Trainer(
+    model=model,
+    args=args,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["validation"],
+    tokenizer=processor,
+    compute_metrics=compute_metrics
+)
+
+# ----------------------------
+# TRAIN
+# ----------------------------
+trainer.train()
+
+# ----------------------------
+# SAVE FINAL MODEL + LABELS
+# ----------------------------
+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)
+
+print(f"✅ Model and processor saved to {OUTPUT_DIR}")