Update main.py

main.py

@@ -50,28 +50,54 @@ except Exception as e:
 caption_model_processor = {"processor": processor, "model": model}
 logger.info("Finished loading models!!!")
 
+# Initialize FastAPI app
 app = FastAPI()
 
+# Define a queue for request processing
+request_queue = asyncio.Queue()
+
+# Define a response model for the processed image
 class ProcessResponse(BaseModel):
     image: str  # Base64 encoded image
     parsed_content_list: str
     label_coordinates: str
 
-# Create a queue for sequential processing
-request_queue = asyncio.Queue()
 
+# Define the async worker function
+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
+        try:
+            await task  # Process the task
+        except Exception as e:
+            logger.error(f"Error while processing task: {e}")
+        finally:
+            request_queue.task_done()  # Mark the task as done
+
+
+# Start the worker when the application starts
+@app.on_event("startup")
+async def startup_event():
+    logger.info("Starting background worker...")
+    asyncio.create_task(worker())  # Start the worker in the background
+
+
+# Define the process 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
         image_save_path = "imgs/saved_image_demo.png"
         os.makedirs(os.path.dirname(image_save_path), exist_ok=True)
-        
-        # Save the image asynchronously
-        buffer = io.BytesIO()
-        image_input.save(buffer, format="PNG")
-        buffer.seek(0)
+
+        # Save the image
+        image_input.save(image_save_path)
+        logger.debug(f"Image saved to: {image_save_path}")
 
         # Perform YOLO and caption model inference
         box_overlay_ratio = image_input.size[0] / 3200
@@ -106,7 +132,7 @@ async def process(image_input: Image.Image, box_threshold: float, iou_threshold:
             iou_threshold=iou_threshold,
         )
 
-        # Convert image to base64
+        # Convert labeled image to base64
         image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
         buffered = io.BytesIO()
         image.save(buffered, format="PNG")
@@ -124,6 +150,8 @@ async def process(image_input: Image.Image, box_threshold: float, iou_threshold:
         logger.error(f"Error in process function: {e}")
         raise
 
+
+# Define the process_image endpoint
 @app.post("/process_image", response_model=ProcessResponse)
 async def process_image(
     image_file: UploadFile = File(...),
@@ -135,16 +163,15 @@ async def process_image(
         contents = await image_file.read()
         image_input = Image.open(io.BytesIO(contents)).convert("RGB")
         
+        # Create a task for processing
+        task = asyncio.create_task(process(image_input, box_threshold, iou_threshold))
+        
         # Add the task to the queue
-        task = asyncio.create_task(
-            process(image_input, box_threshold, iou_threshold)
-        )
         await request_queue.put(task)
+        logger.info(f"Task added to queue. Current queue size: {request_queue.qsize()}")
         
-        # Process the next task in the queue
-        task = await request_queue.get()
+        # Wait for the task to complete
         response = await task
-        request_queue.task_done()
 
         return response
     except Exception as e: