Update Training script

app.py

@@ -16,21 +16,33 @@ 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)
+    try:
+        processor = AutoImageProcessor.from_pretrained(MODEL_DIR)
+        model = BeitForImageClassification.from_pretrained(MODEL_DIR)
+        labels_path = os.path.join(MODEL_DIR, "labels.json")
+        if os.path.exists(labels_path):
+            with open(labels_path) as f:
+                CLASSES = json.load(f)
+        print("✅ Model and processor loaded successfully")
+    except Exception as e:
+        processor, model = None, None
+        print(f"⚠️ Skipping model load: {e}")
 
 @app.on_event("startup")
 def startup_event():
     if os.path.exists(MODEL_DIR):
         load_model()
+    else:
+        print("⚠️ MODEL_DIR not found, skipping model load")
 
 @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")
+    try:
+        img = Image.open(file.file).convert("RGB")
+    except Exception as e:
+        return {"error": f"Invalid image: {str(e)}"}
     inputs = processor(images=img, return_tensors="pt")
     with torch.no_grad():
         logits = model(**inputs).logits
@@ -38,7 +50,7 @@ async def predict(file: UploadFile):
     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)}]
+        "probabilities": {CLASSES[i]: float(p) for i, p in enumerate(probs)}
     }
 
 @app.post("/train")