Update app.py

app.py

@@ -2,45 +2,122 @@ import gradio as gr
 from ultralytics import YOLO
 import spaces
 import torch
+import cv2
+import numpy as np
+import os
+import requests
+
+# Define constants for the new model
+ENTITIES_COLORS = {
+    "Caption": (191, 100, 21),
+    "Footnote": (2, 62, 115),
+    "Formula": (140, 80, 58),
+    "List-item": (168, 181, 69),
+    "Page-footer": (2, 69, 84),
+    "Page-header": (83, 115, 106),
+    "Picture": (255, 72, 88),
+    "Section-header": (0, 204, 192),
+    "Table": (116, 127, 127),
+    "Text": (0, 153, 221),
+    "Title": (196, 51, 2)
+}
+BOX_PADDING = 2
 
 # Load pre-trained YOLOv8 model
-model = YOLO("yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt")
+model_path_1 = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
+model_path_2 = "models/dla-model.pt"
+
+if not os.path.exists(model_path_1):
+    # Download the model file if it doesn't exist
+    model_url_1 = "https://huggingface.co/DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet/resolve/main/yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
+    response = requests.get(model_url_1)
+    with open(model_path_1, "wb") as f:
+        f.write(response.content)
+
+if not os.path.exists(model_path_2):
+    # Assume the second model file is manually uploaded in the specified path
 
-# Get class names from model
-class_names = model.names
+# Load models
+model_1 = YOLO(model_path_1)
+model_2 = YOLO(model_path_2)
+
+# Get class names from the first model
+class_names_1 = model_1.names
+class_names_2 = list(ENTITIES_COLORS.keys())
 
 @spaces.GPU(duration=60)
-def process_image(image):
+def process_image(image, model_choice):
     try:
-        # Process the image
-        results = model(source=image, save=False, show_labels=True, show_conf=True, show_boxes=True)
-        result = results[0]  # Get the first result
+        if model_choice == "YOLOv8 Model":
+            # Use the first model
+            results = model_1(source=image, save=False, show_labels=True, show_conf=True, show_boxes=True)
+            result = results[0]
+
+            # Extract annotated image and labels with class names
+            annotated_image = result.plot()
+
+            detected_areas_labels = "\n".join([
+                f"{class_names_1[int(box.cls.item())].upper()}: {float(box.conf):.2f}" for box in result.boxes
+            ])
 
-        # Extract annotated image and labels with class names
-        annotated_image = result.plot()
+            return annotated_image, detected_areas_labels
+        
+        elif model_choice == "DLA Model":
+            # Use the second model
+            image_path = "input_image.jpg"  # Temporary save the uploaded image
+            cv2.imwrite(image_path, cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR))
+            image = cv2.imread(image_path)
+            results = model_2.predict(source=image, conf=0.2, iou=0.8)
+            boxes = results[0].boxes
 
-        # Use cls attribute for labels and get class name from model
-        detected_areas_labels = "\n".join([
-            f"{class_names[int(box.cls.item())].upper()}: {float(box.conf):.2f}" for box in result.boxes
-        ])
+            if len(boxes) == 0:
+                return image
+
+            for box in boxes:
+                detection_class_conf = round(box.conf.item(), 2)
+                cls = class_names_2[int(box.cls)]
+                start_box = (int(box.xyxy[0][0]), int(box.xyxy[0][1]))
+                end_box = (int(box.xyxy[0][2]), int(box.xyxy[0][3]))
+
+                line_thickness = round(0.002 * (image.shape[0] + image.shape[1]) / 2) + 1
+                image = cv2.rectangle(img=image, 
+                                      pt1=start_box, 
+                                      pt2=end_box,
+                                      color=ENTITIES_COLORS[cls], 
+                                      thickness=line_thickness)
+                
+                text = cls + " " + str(detection_class_conf)
+                font_thickness = max(line_thickness - 1, 1)
+                (text_w, text_h), _ = cv2.getTextSize(text=text, fontFace=2, fontScale=line_thickness/3, thickness=font_thickness)
+                image = cv2.rectangle(img=image,
+                                      pt1=(start_box[0], start_box[1] - text_h - BOX_PADDING*2),
+                                      pt2=(start_box[0] + text_w + BOX_PADDING * 2, start_box[1]),
+                                      color=ENTITIES_COLORS[cls],
+                                      thickness=-1)
+                start_text = (start_box[0] + BOX_PADDING, start_box[1] - BOX_PADDING)
+                image = cv2.putText(img=image, text=text, org=start_text, fontFace=0, color=(255,255,255), fontScale=line_thickness/3, thickness=font_thickness)
+            
+            return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), "Labels: " + ", ".join(class_names_2)
+        
+        else:
+            return None, "Invalid model choice"
 
-        return annotated_image, detected_areas_labels
     except Exception as e:
         return None, f"Error processing image: {e}"
 
 # Create the Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown("# Document Segmentation Demo (ZeroGPU)")
-
+    
     with gr.Row():
+        model_choice = gr.Dropdown(["YOLOv8 Model", "DLA Model"], label="Select Model", value="YOLOv8 Model")
         input_image = gr.Image(type="pil", label="Upload Image")
-        output_image = gr.Image(type="pil", label="Annotated Image")
-
+    
+    output_image = gr.Image(type="pil", label="Annotated Image")
     output_text = gr.Textbox(label="Detected Areas and Labels")
-
-    # Button to trigger inference
+    
     btn = gr.Button("Run Document Segmentation")
-    btn.click(fn=process_image, inputs=input_image, outputs=[output_image, output_text])
+    btn.click(fn=process_image, inputs=[input_image, model_choice], outputs=[output_image, output_text])
 
 # Launch the demo with queuing
 demo.queue(max_size=1).launch()