Update main.py

main.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from pydantic import BaseModel#
+from pydantic import BaseModel
 import base64
 import io
 import os
@@ -24,10 +24,7 @@ yolo_model = get_yolo_model(model_path="weights/best.pt")
 
 # Handle device placement
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-if str(device) == "cuda":
-    yolo_model = yolo_model.cuda()
-else:
-    yolo_model = yolo_model.cpu()
+yolo_model = yolo_model.to(device)
 
 # Load caption model and processor
 try:
@@ -38,7 +35,7 @@ try:
         "weights/icon_caption_florence",
         torch_dtype=torch.float16,
         trust_remote_code=True,
-    ).to("cuda")
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
 except Exception as e:
     logger.warning(f"Failed to load caption model on GPU: {e}. Falling back to CPU.")
     model = AutoModelForCausalLM.from_pretrained(
@@ -48,7 +45,7 @@ except Exception as e:
     )
 
 caption_model_processor = {"processor": processor, "model": model}
-logger.info("Finished loading models!!!")
+logger.info("Finished loading models!")
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -56,51 +53,44 @@ app = FastAPI()
 MAX_QUEUE_SIZE = 10  # Set a reasonable limit based on your system capacity
 request_queue = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)
 
-# Define a response model for the processed image
+# Define response model
 class ProcessResponse(BaseModel):
     image: str  # Base64 encoded image
     parsed_content_list: str
     label_coordinates: str
 
 
-# Define the async worker function
+# Background worker to process queue tasks
 async def worker():
-    """
-    Background worker to process tasks from the request queue sequentially.
-    """
     while True:
-        task = await request_queue.get()  # Get the next task from the queue
+        task = await request_queue.get()
         try:
-            await task  # Process the task
+            await task
         except Exception as e:
             logger.error(f"Error while processing task: {e}")
         finally:
-            request_queue.task_done()  # Mark the task as done
+            request_queue.task_done()
 
 
-# Start the worker when the application starts
+# Start worker on startup
 @app.on_event("startup")
 async def startup_event():
     logger.info("Starting background worker...")
-    
-    asyncio.create_task(worker())  # Start the worker in the background
+    asyncio.create_task(worker())
 
 
-# Define the process function
+# Image processing function
 async def process(image_input: Image.Image, box_threshold: float, iou_threshold: float) -> ProcessResponse:
-    """
-    Asynchronously processes an image using YOLO and caption models.
-    """
     try:
-        # Define the save path and ensure the directory exists
+        # Define save path
         image_save_path = "imgs/saved_image_demo.png"
         os.makedirs(os.path.dirname(image_save_path), exist_ok=True)
 
-        # Save the image
+        # Save image
         image_input.save(image_save_path)
         logger.debug(f"Image saved to: {image_save_path}")
 
-        # Perform YOLO and caption model inference
+        # YOLO and caption model inference
         box_overlay_ratio = image_input.size[0] / 3200
         draw_bbox_config = {
             "text_scale": 0.8 * box_overlay_ratio,
@@ -152,7 +142,7 @@ async def process(image_input: Image.Image, box_threshold: float, iou_threshold:
         raise HTTPException(status_code=500, detail=f"Failed to process the image: {e}")
 
 
-# Define the process_image endpoint
+# API endpoint for processing images
 @app.post("/process_image", response_model=ProcessResponse)
 async def process_image(
     image_file: UploadFile = File(...),
@@ -160,22 +150,22 @@ async def process_image(
     iou_threshold: float = 0.1,
 ):
     try:
-        # Read the image file
+        # Read image file
         contents = await image_file.read()
         try:
             image_input = Image.open(io.BytesIO(contents)).convert("RGB")
-        except UnidentifiedImageError as e:
-            logger.error(f"Unsupported image format: {e}")
+        except UnidentifiedImageError:
+            logger.error("Unsupported image format.")
             raise HTTPException(status_code=400, detail="Unsupported image format.")
 
-        # Create a task for processing
+        # Create processing task
         task = asyncio.create_task(process(image_input, box_threshold, iou_threshold))
 
-        # Add the task to the queue
+        # Add task to 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
+        # Wait for processing to complete
         response = await task
 
         return response
@@ -183,4 +173,4 @@ async def process_image(
         raise he
     except Exception as e:
         logger.error(f"Error processing image: {e}")
-        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")#
+        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")