Update app.py

app.py

@@ -6,14 +6,19 @@ from pydantic import BaseModel
 from io import BytesIO
 from PIL import Image
 
+# Initialize the FastAPI app
 app = FastAPI()
 
-# Load the pretrained MobileNetV2 model from torchvision
-model = models.mobilenet_v2(pretrained=True)
-model.classifier[1] = nn.Linear(model.last_channel, 3)  # Replace with the number of classes in your task
+# Path to the trained model checkpoint
+MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/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
 model.eval()
 
-# Define the image preprocessing transform
+# Define image preprocessing transformations
 preprocess = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -24,7 +29,7 @@ preprocess = transforms.Compose([
 def preprocess_image(image: Image.Image):
     image = image.convert("RGB")
     image = preprocess(image)
-    image = image.unsqueeze(0)
+    image = image.unsqueeze(0)  # Add a batch dimension
     return image
 
 @app.post("/predict/")
@@ -40,7 +45,7 @@ async def predict(file: UploadFile = File(...)):
             outputs = model(image_tensor)
             predicted_class_idx = torch.argmax(outputs, dim=1).item()
 
-        # Map the predicted class index to the class labels
+        # Map class index to class label
         class_labels = ["Cocci", "Bacilli", "Spirilla"]  # Replace with your actual class labels
         predicted_class = class_labels[predicted_class_idx]
 
@@ -49,4 +54,5 @@ async def predict(file: UploadFile = File(...)):
     except Exception as e:
         return {"error": str(e)}
 
-# Run the app with: uvicorn app:app --reload
+# To run the app, use the following command in the terminal:
+# uvicorn app:app --host 0.0.0.0 --port 8000 --reload