Update main.py

main.py

@@ -1,5 +1,4 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 import base64
 import io
@@ -17,31 +16,20 @@ from utils import (
     get_som_labeled_img,
 )
 
-# Import YOLO from ultralytics and transformers for captioning
+# Load the YOLO model using the ultralytics class instead of torch.load
 from ultralytics import YOLO
-from transformers import AutoProcessor, AutoModelForCausalLM
-
-# ---------------------------------------------------------------------------
-# Load the YOLO model
-# ---------------------------------------------------------------------------
-try:
-    yolo_model = torch.load("weights/icon_detect/best.pt", map_location="cuda", weights_only=False)["model"]
-    yolo_model = yolo_model.to("cuda")
-except Exception as e:
-    print("Error loading YOLO model on CUDA:", e)
-    yolo_model = torch.load("weights/icon_detect/best.pt", map_location="cpu", weights_only=False)["model"]
 
+# Use the YOLO constructor to load the model properly
+yolo_model = YOLO("weights/icon_detect/best.pt")
 print(f"YOLO model type: {type(yolo_model)}")
 
-# ---------------------------------------------------------------------------
 # Load the captioning model (Florence-2)
-# ---------------------------------------------------------------------------
+from transformers import AutoProcessor, AutoModelForCausalLM
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
-# Load the processor for the Florence-2 model
 processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base", trust_remote_code=True)
-
 try:
     model = AutoModelForCausalLM.from_pretrained(
         "weights/icon_caption_florence",
@@ -50,14 +38,12 @@ try:
     ).to(device)
 except Exception as e:
     print(f"Error loading caption model: {str(e)}")
-    # Fallback to CPU with float32
     model = AutoModelForCausalLM.from_pretrained(
         "weights/icon_caption_florence",
         torch_dtype=torch.float32,
         trust_remote_code=True
     ).to("cpu")
 
-# Force configuration for DaViT vision tower if missing
 if not hasattr(model.config, 'vision_config'):
     model.config.vision_config = {}
 if 'model_type' not in model.config.vision_config:
@@ -66,9 +52,6 @@ if 'model_type' not in model.config.vision_config:
 caption_model_processor = {"processor": processor, "model": model}
 print("Finish loading caption model!")
 
-# ---------------------------------------------------------------------------
-# Create FastAPI application and response model
-# ---------------------------------------------------------------------------
 app = FastAPI()
 
 class ProcessResponse(BaseModel):
@@ -76,18 +59,13 @@ class ProcessResponse(BaseModel):
     parsed_content_list: str
     label_coordinates: str
 
-# ---------------------------------------------------------------------------
-# Main processing function
-# ---------------------------------------------------------------------------
 def process(image_input: Image.Image, box_threshold: float, iou_threshold: float) -> ProcessResponse:
-    # Save the input image temporarily
     image_save_path = "imgs/saved_image_demo.png"
     os.makedirs(os.path.dirname(image_save_path), exist_ok=True)
     image_input.save(image_save_path)
     
-    # Open the saved image for processing
     image = Image.open(image_save_path)
-    box_overlay_ratio = image.size[0] / 3200  # adjust scaling factor as needed
+    box_overlay_ratio = image.size[0] / 3200
     draw_bbox_config = {
         "text_scale": 0.8 * box_overlay_ratio,
         "text_thickness": max(int(2 * box_overlay_ratio), 1),
@@ -95,7 +73,6 @@ def process(image_input: Image.Image, box_threshold: float, iou_threshold: float
         "thickness": max(int(3 * box_overlay_ratio), 1),
     }
 
-    # Run OCR to get text and OCR bounding boxes
     ocr_bbox_rslt, is_goal_filtered = check_ocr_box(
         image_save_path,
         display_img=False,
@@ -106,7 +83,6 @@ def process(image_input: Image.Image, box_threshold: float, iou_threshold: float
     )
     text, ocr_bbox = ocr_bbox_rslt
 
-    # Run YOLO and semantic processing to get the labeled image and bounding boxes
     dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(
         image_save_path,
         yolo_model,
@@ -118,13 +94,10 @@ def process(image_input: Image.Image, box_threshold: float, iou_threshold: float
         ocr_text=text,
         iou_threshold=iou_threshold,
     )
-    
-    # Decode the base64-encoded image output from get_som_labeled_img
     image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
     print("Finish processing")
     parsed_content_list_str = "\n".join(parsed_content_list)
 
-    # Encode final image to base64 string for response
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
     img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
@@ -135,9 +108,6 @@ def process(image_input: Image.Image, box_threshold: float, iou_threshold: float
         label_coordinates=str(label_coordinates),
     )
 
-# ---------------------------------------------------------------------------
-# FastAPI endpoint for image processing
-# ---------------------------------------------------------------------------
 @app.post("/process_image", response_model=ProcessResponse)
 async def process_image(
     image_file: UploadFile = File(...),
@@ -148,13 +118,10 @@ async def process_image(
         contents = await image_file.read()
         image_input = Image.open(io.BytesIO(contents)).convert("RGB")
         
-        # Debug logging for file information
         print(f"Processing image: {image_file.filename}")
         print(f"Image size: {image_input.size}")
         
         response = process(image_input, box_threshold, iou_threshold)
-        
-        # Validate response
         if not response.image:
             raise ValueError("Empty image in response")
             
@@ -162,5 +129,5 @@ async def process_image(
         
     except Exception as e:
         import traceback
-        traceback.print_exc()  # Print full traceback for debugging
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))