Update main.py

main.py

@@ -4,7 +4,7 @@ import base64
 import io
 import os
 import logging
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 import torch
 import asyncio
 from utils import (
@@ -149,7 +149,7 @@ async def process(image_input: Image.Image, box_threshold: float, iou_threshold:
         )
     except Exception as e:
         logger.error(f"Error in process function: {e}")
-        raise
+        raise HTTPException(status_code=500, detail=f"Failed to process the image: {e}")
 
 
 # Define the process_image endpoint
@@ -162,19 +162,25 @@ async def process_image(
     try:
         # Read the image file
         contents = await image_file.read()
-        image_input = Image.open(io.BytesIO(contents)).convert("RGB")
-        
+        try:
+            image_input = Image.open(io.BytesIO(contents)).convert("RGB")
+        except UnidentifiedImageError as e:
+            logger.error(f"Unsupported image format: {e}")
+            raise HTTPException(status_code=400, detail="Unsupported image format.")
+
         # Create a task for processing
         task = asyncio.create_task(process(image_input, box_threshold, iou_threshold))
-        
+
         # Add the task to the queue
         await request_queue.put(task)
         logger.info(f"Task added to queue. Current queue size: {request_queue.qsize()}")
-        
+
         # Wait for the task to complete
         response = await task
 
         return response
+    except HTTPException as he:
+        raise he
     except Exception as e:
         logger.error(f"Error processing image: {e}")
-        raise HTTPException(status_code=500, detail=str(e))#
+        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")