Update app.py

app.py

@@ -1,33 +1,78 @@
+from flask import Flask, request, jsonify
 import torch
 import torch.nn as nn
-import torch.hub
-import requests
-from torchvision import models, transforms
+from torchvision import transforms
 from PIL import Image
-from fastapi import FastAPI, UploadFile, File
-from io import BytesIO
+import io
 
-app = FastAPI()
+# Define the model architecture that matches the saved .pth file
+class BacterialMorphologyClassifier(nn.Module):
+    def __init__(self):
+        super(BacterialMorphologyClassifier, self).__init__()
+        self.feature_extractor = nn.Sequential(
+            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        )
+        self.fc = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(64 * 56 * 56, 128),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(128, 3),
+            nn.Softmax(dim=1),
+        )
+    
+    def forward(self, x):
+        x = self.feature_extractor(x)
+        x = self.fc(x)
+        return x
 
-# Load the model
+# Load the model and weights
+model = BacterialMorphologyClassifier()
 MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/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')))
+state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
+model.load_state_dict(state_dict)
 model.eval()
 
-# Define image transformation
+# Set up Flask app
+app = Flask(__name__)
+
+# Define image preprocessing transformations
 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.post("/predict/")
-async def predict(file: UploadFile = File(...)):
-    image = Image.open(BytesIO(await file.read())).convert("RGB")
-    image_tensor = transform(image).unsqueeze(0)
-    with torch.no_grad():
+@app.route('/predict', methods=['POST'])
+def predict():
+    try:
+        # Get image from request
+        image_file = request.files['image']
+        image = Image.open(io.BytesIO(image_file.read())).convert('RGB')
+        
+        # Preprocess the image
+        image_tensor = transform(image).unsqueeze(0)
+        
+        # Make prediction
         output = model(image_tensor)
-        _, predicted = output.max(1)
-        return {"predicted_class": int(predicted.item())}
+        prediction = output.argmax().item()
+        
+        # Class mapping
+        class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
+        
+        # Return prediction result
+        response = {
+            'predicted_class': class_labels[prediction],
+            'confidence': output.max().item()
+        }
+        return jsonify(response)
+    except Exception as e:
+        return jsonify({'error': str(e)})
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000, debug=True)