Update app.py

app.py

@@ -23,52 +23,54 @@ ENTITIES_COLORS = {
 }
 BOX_PADDING = 2
 
-# Load pre-trained YOLOv8 model
-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)
+# Load pre-trained YOLOv8 models
+model_paths = {
+    "YOLOv8x Model": "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt",
+    "YOLOv8m Model": "yolov8m-doclaynet.pt",
+    "YOLOv8n Model": "yolov8n-doclaynet.pt",
+    "DLA Model": "models/dla-model.pt"
+}
 
-if not os.path.exists(model_path_2):
-    # Assume the second model file is manually uploaded in the specified path
-    pass
+# Ensure the model files are in the correct location
+for model_name, model_path in model_paths.items():
+    if not os.path.exists(model_path):
+        # For demonstration, we only download the YOLOv8x model
+        if model_name == "YOLOv8x Model":
+            model_url = "https://huggingface.co/DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet/resolve/main/yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
+            response = requests.get(model_url)
+            with open(model_path, "wb") as f:
+                f.write(response.content)
 
 # Load models
-model_1 = YOLO(model_path_1)
-model_2 = YOLO(model_path_2)
+models = {name: YOLO(path) for name, path in model_paths.items()}
 
-# Get class names from the first model
-class_names_1 = model_1.names
-class_names_2 = list(ENTITIES_COLORS.keys())
+# Get class names from the YOLOv8 models
+class_names = list(ENTITIES_COLORS.keys())
 
 @spaces.GPU(duration=60)
 def process_image(image, model_choice):
     try:
-        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)
+        if "YOLOv8" in model_choice:
+            # Use the selected YOLOv8 model
+            model = models[model_choice]
+            results = model(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
+                f"{class_names[int(box.cls.item())].upper()}: {float(box.conf):.2f}" for box in result.boxes
             ])
 
             return annotated_image, detected_areas_labels
         
         elif model_choice == "DLA Model":
-            # Use the second model
+            # Use the DLA 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)
+            results = models[model_choice].predict(source=image, conf=0.2, iou=0.8)
             boxes = results[0].boxes
 
             if len(boxes) == 0:
@@ -76,7 +78,7 @@ def process_image(image, model_choice):
 
             for box in boxes:
                 detection_class_conf = round(box.conf.item(), 2)
-                cls = class_names_2[int(box.cls)]
+                cls = class_names[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]))
 
@@ -98,7 +100,7 @@ def process_image(image, model_choice):
                 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)
+            return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), "Labels: " + ", ".join(class_names)
         
         else:
             return None, "Invalid model choice"
@@ -114,7 +116,7 @@ with gr.Blocks() as demo:
         input_image = gr.Image(type="pil", label="Upload Image")
         output_image = gr.Image(type="pil", label="Annotated Image")
 
-    model_choice = gr.Dropdown(["YOLOv8 Model", "DLA Model"], label="Select Model", value="YOLOv8 Model", scale=0.5)
+    model_choice = gr.Dropdown(list(model_paths.keys()), label="Select Model", value="YOLOv8x Model", scale=0.5)
     output_text = gr.Textbox(label="Detected Areas and Labels")
     
     btn = gr.Button("Run Document Segmentation")