api code changed

src/serve.py

@@ -1,68 +1,54 @@
+from fastapi import FastAPI, HTTPException, File, UploadFile
+from fastapi.responses import JSONResponse
+from io import BytesIO
 from PIL import Image
-from torchvision import transforms
 import torch
+from torchvision import transforms
 import os
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from io import BytesIO
-
-# from dotenv import load_dotenv
-from .model import MalwareNet, malware_classes  # assuming malware_classes contains class names
-
-# load_dotenv()
+from .model import MalwareNet, malware_classes
 
 app = FastAPI()
 
-# Preprocessing function for the model
-def preprocess_image(image_path):
-    image = Image.open(image_path).convert("RGB")
+def preprocess_image(image_data):
+    image = Image.open(BytesIO(image_data)).convert("RGB")
     preprocess = transforms.Compose([
-        transforms.Resize((224, 224)),  # Resize to the input size expected by the model
-        transforms.ToTensor(),          # Convert to tensor
-        transforms.Normalize(           # Normalize using model's requirements (e.g. ImageNet)
-            mean=[0.485, 0.456, 0.406], 
-            std=[0.229, 0.224, 0.225]
-        )
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])
-    return preprocess(image).unsqueeze(0)  # Add batch dimension
+    return preprocess(image).unsqueeze(0)
 
-# Load model and its weights
 def load_model():
     model = MalwareNet()
     base_dir = os.path.dirname(os.path.abspath(__file__))
-    model_location = os.path.join(base_dir, '../model/malwareNet.pt')  # Relative path to the model file
+    model_location = os.path.join(base_dir, '../model/malwareNet.pt')
     state_dict = torch.load(model_location, map_location=torch.device('cpu'), weights_only=True)
     model.load_state_dict(state_dict)
-    model.eval()  # Set the model to evaluation mode
+    model.eval()
     return model
 
-@app.get("/")
-def status():
-    return {"status": "ok"}
-
 @app.post("/predict")
-async def predict(data: dict):
-    image_path = data.get("image_url")
-    if not os.path.exists(image_path):
-        raise HTTPException(status_code=400, detail="Image path does not exist.")
-
+async def predict(file: UploadFile = File(...)):
     try:
-        # Load and preprocess the image
-        img_tensor = preprocess_image(image_path)
+        # Read file bytes
+        image_data = await file.read()
+
+        # Preprocess the image
+        img_tensor = preprocess_image(image_data)
 
         # Load the model and make the prediction
         model = load_model()
-        with torch.no_grad():  # No gradient calculation is needed
+        with torch.no_grad():
             prediction = model(img_tensor)
 
         # Get the predicted class
         predicted_class = malware_classes[torch.argmax(prediction).item()]
 
-        return JSONResponse(content={"image": image_path, "prediction": predicted_class})
-
+        return JSONResponse(content={"prediction": predicted_class})
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error processing the image: {e}")
 
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(