try new version of API

main.py

@@ -14,25 +14,7 @@ import torch
 import logging
 from typing import List
 import httpx
-from concurrent.futures import ProcessPoolExecutor
-
-
-def process_single_image(image_url, model):
-    try:
-        response = requests.get(image_url)
-        image = Image.open(BytesIO(response.content))
-        processed_image = process_image(image, size=image_size)
-        image_tensor = transforms.ToTensor()(processed_image).unsqueeze(0)
-
-        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)
-        return {"imageUrl": image_url, "confidence": confidence}
-    except Exception as e:
-        return {"imageUrl": image_url, "error": str(e)}
-
+import asyncio
 
 app = FastAPI()
 
@@ -72,6 +54,8 @@ class PredictRequest(BaseModel):
     modelName: str
 
 
+torch.set_num_threads(6)
+
 # Dictionnaire pour stocker les pipelines de modèles
 model_pipelines = {}
 
@@ -144,61 +128,86 @@ class BatchPredictRequest(BaseModel):
     modelName: str
 
 
-# @app.post("/batch_predict")
-# async def batch_predict(request: BatchPredictRequest):
-#     model_name = request.modelName
-#     results = []
+@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")
+    # 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]
+    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
+    # 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)
+            # Preprocess the image
+            processed_image = process_image(image, size=image_size)
 
-#             # Convert to tensor
-#             image_tensor = transforms.ToTensor()(processed_image).unsqueeze(0)
+            # 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)
+            # 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})
+            results.append({"imageUrl": image_url, "confidence": confidence})
 
-#     # Return the results as JSON
-#     return JSONResponse(content={"results": results})
+    # Return the results as JSON
+    return JSONResponse(content={"results": results})
 
 
 @app.post("/batch_predict")
 async def batch_predict(request: BatchPredictRequest):
     model_name = request.modelName
+
+    # 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]
-
-    with ProcessPoolExecutor() as executor:
-        results = list(
-            executor.map(
-                process_single_image,
-                request.imageUrls,
-                [model] * len(request.imageUrls),
-            )
-        )
-
+    semaphore = asyncio.Semaphore(
+        6
+    )  # Limiter à 8 tâches simultanées pour éviter de surcharger la machine
+
+    async def process_single_image(image_url):
+        async with semaphore:
+            try:
+                async with httpx.AsyncClient() as client:
+                    response = await client.get(image_url)
+                    image = Image.open(BytesIO(response.content))
+            except Exception:
+                return {"imageUrl": image_url, "error": "Invalid image URL"}
+
+            # 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)
+
+            return {"imageUrl": image_url, "confidence": confidence}
+
+    # Launch tasks in parallel
+    tasks = [process_single_image(url) for url in request.imageUrls]
+    results = await asyncio.gather(*tasks)
+
+    # Return the results as JSON
     return JSONResponse(content={"results": results})