Update main.py

main.py

@@ -4,57 +4,58 @@ from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
 import io
 from torchvision import transforms
-
-# Import the loader from the file next to this one
 from model_loader import load_model
 
 app = FastAPI()
 
-# Enable CORS so React can talk to this
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], 
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+model = None
+device = torch.device("cpu")
+
 # --- LOAD MODEL ---
-print("Loading model...")
+print("--- STARTING SERVER ---")
 try:
-    # This expects the .pth file to be in the root folder
-    model_wrapper = load_model("InceptionViT_best_model.pth")
-    model = model_wrapper.model 
-    print("Model loaded successfully!")
+    model = load_model("InceptionViT_best_model.pth")
+    print("✅ Model loaded successfully!")
 except Exception as e:
-    print(f"CRITICAL ERROR LOADING MODEL: {e}")
-    # If this fails, check the troubleshooting note below
+    print(f"❌ CRITICAL ERROR: {e}")
 
 # --- TRANSFORM ---
-# Ensure these match your training exactly
+# Matches your training code exactly
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
 ])
 
 @app.get("/")
 def home():
-    return {"message": "API is running"}
+    return {"status": "Running"}
 
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
-    # 1. Read Image
+    if model is None:
+        return {"error": "Model not loaded"}
+    
     image_data = await file.read()
     image = Image.open(io.BytesIO(image_data)).convert("RGB")
     
-    # 2. Preprocess
-    tensor = transform(image).unsqueeze(0) # Add batch dimension
+    tensor = transform(image).unsqueeze(0).to(device)
     
-    # 3. Predict
     with torch.no_grad():
-        outputs = model(tensor)
-        _, predicted = torch.max(outputs, 1)
+        logits = model(tensor)
+        probabilities = torch.nn.functional.softmax(logits, dim=1)
+        confidence, predicted = torch.max(probabilities, 1)
         
-    # 4. Return
-    return {"result": str(predicted.item())}
+    return {
+        "prediction": str(predicted.item()),
+        "confidence": float(confidence.item())
+    }