add new route

Dockerfile

@@ -30,4 +30,4 @@ EXPOSE 7860
 #   git clone $(cat /run/secrets/api_read)
 
 # Commande pour lancer l'application
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860",  "--workers", "4"]
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -12,6 +12,9 @@ from huggingface_hub import hf_hub_download
 from architecture.resnet import ResNet
 import torch
 import logging
+from typing import List
+import httpx
+
 
 app = FastAPI()
 
@@ -116,3 +119,48 @@ async def predict(request: PredictRequest):
     logging.info("confidence: %s", confidence)
     # Return the probabilities as JSON
     return JSONResponse(content={"confidence": confidence})
+
+
+class BatchPredictRequest(BaseModel):
+    imageUrls: List[str]
+    modelName: str
+
+
+@app.post("/batch_predict")
+async def batch_predict(request: BatchPredictRequest):
+    model_name = request.modelName
+    results = []
+
+    # Verify if the model is loaded
+    if model_name not in model_pipelines:
+        raise HTTPException(status_code=404, detail="Model not found")
+
+    model = model_pipelines[model_name]
+
+    # Asynchronously process each image
+    async with httpx.AsyncClient() as client:
+        for image_url in request.imageUrls:
+            try:
+                response = await client.get(image_url)
+                image = Image.open(BytesIO(response.content))
+            except Exception as e:
+                results.append({"imageUrl": image_url, "error": "Invalid image URL"})
+                continue
+
+            # Preprocess the image
+            processed_image = process_image(image, size=image_size)
+
+            # Convert to tensor
+            image_tensor = transforms.ToTensor()(processed_image).unsqueeze(0)
+
+            # Perform inference
+            with torch.no_grad():
+                outputs = model(image_tensor)
+                probabilities = torch.nn.functional.softmax(outputs, dim=1)
+                predicted_probabilities = probabilities.numpy().tolist()
+                confidence = round(predicted_probabilities[0][1], 2)
+
+            results.append({"imageUrl": image_url, "confidence": confidence})
+
+    # Return the results as JSON
+    return JSONResponse(content={"results": results})

requirements.txt

@@ -6,4 +6,5 @@ requests
 torchvision
 huggingface_hub
 torch
-numpy
+numpy
+httpx