Update app.py

app.py

@@ -1,58 +1,33 @@
 import torch
 import torch.nn as nn
+import torch.hub
+import requests
 from torchvision import models, transforms
+from PIL import Image
 from fastapi import FastAPI, UploadFile, File
-from pydantic import BaseModel
 from io import BytesIO
-from PIL import Image
 
-# Initialize the FastAPI app
 app = FastAPI()
 
-# Path to the trained model checkpoint
-MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/blob/main/BacterialMorphologyClassification_model.ipynb"
-
-# Load the pretrained MobileNetV2 model and modify the classifier
-model = models.mobilenet_v2(weights=None)  # No initial weights
-model.classifier[1] = nn.Linear(model.last_channel, 3)  # Replace with 3 classes for your dataset
-model.load_state_dict(torch.load(MODEL_PATH))  # Load your trained model weights
+# Load the model
+MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/BacterialMorphologyClassification_model.pth"
+model = models.mobilenet_v2(weights=None)
+model.classifier[1] = nn.Linear(model.last_channel, 3)
+model.load_state_dict(torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu')))
 model.eval()
 
-# Define image preprocessing transformations
-preprocess = transforms.Compose([
+# Define image transformation
+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])
 ])
 
-# Helper function to preprocess the image
-def preprocess_image(image: Image.Image):
-    image = image.convert("RGB")
-    image = preprocess(image)
-    image = image.unsqueeze(0)  # Add a batch dimension
-    return image
-
 @app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
-    try:
-        # Read and preprocess the image
-        image_data = await file.read()
-        image = Image.open(BytesIO(image_data))
-        image_tensor = preprocess_image(image)
-
-        # Perform inference
-        with torch.no_grad():
-            outputs = model(image_tensor)
-            predicted_class_idx = torch.argmax(outputs, dim=1).item()
-
-        # Map class index to class label
-        class_labels = ["Cocci", "Bacilli", "Spirilla"]  # Replace with your actual class labels
-        predicted_class = class_labels[predicted_class_idx]
-
-        return {"predicted_class": predicted_class}
-
-    except Exception as e:
-        return {"error": str(e)}
-
-# To run the app, use the following command in the terminal:
-# uvicorn app:app --host 0.0.0.0 --port 8000 --reload
+    image = Image.open(BytesIO(await file.read())).convert("RGB")
+    image_tensor = transform(image).unsqueeze(0)
+    with torch.no_grad():
+        output = model(image_tensor)
+        _, predicted = output.max(1)
+        return {"predicted_class": int(predicted.item())}