Update app.py

app.py

@@ -1,49 +1,44 @@
 import torch
+import torch.nn as nn
+from torchvision import models, transforms
 from fastapi import FastAPI, UploadFile, File
 from pydantic import BaseModel
 from io import BytesIO
 from PIL import Image
-from torchvision import transforms
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 
-# Initialize the FastAPI app
 app = FastAPI()
 
-# Load the model and feature extractor
-model = AutoModelForImageClassification.from_pretrained("yolac/BacterialMorphologyClassification")
-feature_extractor = AutoFeatureExtractor.from_pretrained("yolac/BacterialMorphologyClassification")
+# 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
+model.eval()
 
 # Define the image preprocessing transform
 preprocess = transforms.Compose([
-    transforms.Resize((224, 224)),  # Adjust size as needed for your model
+    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")  # Ensure the image is in RGB format
+    image = image.convert("RGB")
     image = preprocess(image)
-    image = image.unsqueeze(0)  # Add a batch dimension
+    image = image.unsqueeze(0)
     return image
 
-# Define the prediction endpoint
 @app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
     try:
-        # Read the uploaded image
+        # Read and preprocess the image
         image_data = await file.read()
         image = Image.open(BytesIO(image_data))
-
-        # Preprocess the image
         image_tensor = preprocess_image(image)
 
         # Perform inference
-        model.eval()
         with torch.no_grad():
             outputs = model(image_tensor)
-            logits = outputs.logits  # Get the model's raw output
-            predicted_class_idx = torch.argmax(logits, dim=1).item()
+            predicted_class_idx = torch.argmax(outputs, dim=1).item()
 
         # Map the predicted class index to the class labels
         class_labels = ["Cocci", "Bacilli", "Spirilla"]  # Replace with your actual class labels
@@ -54,4 +49,4 @@ async def predict(file: UploadFile = File(...)):
     except Exception as e:
         return {"error": str(e)}
 
-# Run the app using the command: uvicorn app:app --reload
+# Run the app with: uvicorn app:app --reload