Update app.py

app.py

@@ -3,6 +3,7 @@ os.environ["GRADIO_TEMP_DIR"] = "./tmp"
 
 import sys
 import torch
+import torchvision
 import gradio as gr
 import numpy as np
 import cv2
@@ -13,10 +14,10 @@ from transformers import (
     RTDetrImageProcessor,
 )
 
-# == select device ==
+# == Device configuration ==
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-# Available models with their corresponding model classes
+# == Model configurations ==
 MODELS = {
     "Egret XLarge": {
         "path": "ds4sd/docling-layout-egret-xlarge",
@@ -40,34 +41,22 @@ MODELS = {
     }
 }
 
-# Classes mapping for the docling model
+# == Class mappings ==
 classes_map = {
-    0: "Caption",
-    1: "Footnote", 
-    2: "Formula",
-    3: "List-item",
-    4: "Page-footer",
-    5: "Page-header",
-    6: "Picture",
-    7: "Section-header",
-    8: "Table",
-    9: "Text",
-    10: "Title",
-    11: "Document Index",
-    12: "Code",
-    13: "Checkbox-Selected",
-    14: "Checkbox-Unselected",
-    15: "Form",
-    16: "Key-Value Region",
+    0: "Caption", 1: "Footnote", 2: "Formula", 3: "List-item",
+    4: "Page-footer", 5: "Page-header", 6: "Picture", 7: "Section-header",
+    8: "Table", 9: "Text", 10: "Title", 11: "Document Index",
+    12: "Code", 13: "Checkbox-Selected", 14: "Checkbox-Unselected", 
+    15: "Form", 16: "Key-Value Region",
 }
 
-# Global variables for model
+# == Global model variables ==
 current_model = None
 current_processor = None
 current_model_name = None
 
 def colormap(N=256, normalized=False):
-    """Generate the color map."""
+    """Generate dynamic colormap."""
     def bitget(byteval, idx):
         return ((byteval & (1 << idx)) != 0)
 
@@ -84,25 +73,24 @@ def colormap(N=256, normalized=False):
     
     if normalized:
         cmap = cmap.astype(np.float32) / 255.0
-
     return cmap
 
 def iomin(box1, box2):
-    """Intersection over Minimum (IoMin)"""
+    """Intersection over Minimum (IoMin)."""
     x1 = torch.max(box1[:, 0], box2[:, 0])
     y1 = torch.max(box1[:, 1], box2[:, 1])
     x2 = torch.min(box1[:, 2], box2[:, 2])
     y2 = torch.min(box1[:, 3], box2[:, 3])
     inter_area = torch.clamp(x2 - x1, min=0) * torch.clamp(y2 - y1, min=0)
-
+    
     box1_area = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])
     box2_area = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])
     min_area = torch.min(box1_area, box2_area)
-
+    
     return inter_area / min_area
 
-def nms(boxes, scores, iou_threshold=0.5):
-    """Custom NMS implementation using IoMin"""
+def nms_custom(boxes, scores, iou_threshold=0.5):
+    """Custom NMS implementation using IoMin."""
     keep = []
     _, order = scores.sort(descending=True)
 
@@ -123,18 +111,19 @@ def nms(boxes, scores, iou_threshold=0.5):
     return torch.tensor(keep, dtype=torch.long)
 
 def load_model(model_name):
-    """Load the selected model"""
+    """Load the selected model."""
     global current_model, current_processor, current_model_name
     
     if current_model_name == model_name:
         return f"✅ Model {model_name} is already loaded!"
     
     try:
-        print(f"Loading model: {model_name}")
         model_info = MODELS[model_name]
         model_path = model_info["path"]
         model_class = model_info["model_class"]
         
+        print(f"Loading {model_name} from {model_path}")
+        
         processor = RTDetrImageProcessor.from_pretrained(model_path)
         model = model_class.from_pretrained(model_path)
         model = model.to(device)
@@ -147,10 +136,11 @@ def load_model(model_name):
         return f"✅ Successfully loaded {model_name}!"
         
     except Exception as e:
+        print(f"Error loading model: {e}")
         return f"❌ Error loading {model_name}: {str(e)}"
 
 def visualize_bbox(image_input, bboxes, classes, scores, id_to_names, alpha=0.3):
-    """Visualize bounding boxes with transparent overlays using OpenCV"""
+    """Visualize bounding boxes with OpenCV."""
     if isinstance(image_input, Image.Image):
         image = np.array(image_input)
         image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
@@ -162,12 +152,12 @@ def visualize_bbox(image_input, bboxes, classes, scores, id_to_names, alpha=0.3)
     else:
         raise ValueError("Input must be PIL Image or numpy array")
 
-    overlay = image.copy()
-    cmap = colormap(N=len(id_to_names), normalized=False)
-
     if len(bboxes) == 0:
         return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
+    overlay = image.copy()
+    cmap = colormap(N=len(id_to_names), normalized=False)
+
     for i in range(len(bboxes)):
         try:
             bbox = bboxes[i]
@@ -186,43 +176,52 @@ def visualize_bbox(image_input, bboxes, classes, scores, id_to_names, alpha=0.3)
             class_id = int(class_id)
             class_name = id_to_names.get(class_id, f"unknown_{class_id}")
 
-            text = f"{class_name}:{score:.3f}"
+            text = f"{class_name}: {score:.3f}"
             color = tuple(int(c) for c in cmap[class_id % len(cmap)])
 
+            # Draw filled rectangle on overlay
             cv2.rectangle(overlay, (x_min, y_min), (x_max, y_max), color, -1)
-            cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, 2)
+            # Draw border on main image
+            cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, 3)
 
-            (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
-            cv2.rectangle(image, (x_min, y_min - text_height - baseline), (x_min + text_width, y_min), color, -1)
-            cv2.putText(image, text, (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+            # Add text label
+            (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)
+            cv2.rectangle(image, (x_min, y_min - text_height - baseline - 4), 
+                         (x_min + text_width + 8, y_min), color, -1)
+            cv2.putText(image, text, (x_min + 4, y_min - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
 
         except Exception as e:
             print(f"Skipping box {i} due to error: {e}")
 
+    # Apply transparency
     cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
+    
     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
-def recognize_image(input_img, conf_threshold, iou_threshold, nms_method, alpha):
-    """Process image with docling layout model"""
+def process_image(input_img, conf_threshold, iou_threshold, nms_method, alpha):
+    """Process image with document layout detection."""
     if input_img is None:
-        return None, "Please upload an image first."
+        return None, "❌ Please upload an image first."
         
     if current_model is None or current_processor is None:
-        return None, "Please load a model first."
+        return None, "❌ Please load a model first."
         
     try:
+        # Prepare image
         if isinstance(input_img, np.ndarray):
             input_img = Image.fromarray(input_img)
         
         if input_img.mode != 'RGB':
             input_img = input_img.convert('RGB')
         
+        # Process with model
         inputs = current_processor(images=[input_img], return_tensors="pt")
         inputs = {k: v.to(device) for k, v in inputs.items()}
         
         with torch.no_grad():
             outputs = current_model(**inputs)
             
+        # Post-process results
         results = current_processor.post_process_object_detection(
             outputs,
             target_sizes=torch.tensor([input_img.size[::-1]]),
@@ -230,7 +229,7 @@ def recognize_image(input_img, conf_threshold, iou_threshold, nms_method, alpha)
         )
         
         if not results or len(results) == 0:
-            return np.array(input_img), "No detections found."
+            return np.array(input_img), "ℹ️ No detections found."
             
         result = results[0]
         boxes = result["boxes"]
@@ -238,116 +237,241 @@ def recognize_image(input_img, conf_threshold, iou_threshold, nms_method, alpha)
         labels = result["labels"]
         
         if len(boxes) == 0:
-            return np.array(input_img), "No detections above confidence threshold."
+            return np.array(input_img), f"ℹ️ No detections above threshold {conf_threshold:.2f}."
         
+        # Apply NMS
         if iou_threshold < 1.0:
             if nms_method == "Custom IoMin":
-                keep_indices = nms(boxes=boxes, scores=scores, iou_threshold=iou_threshold)
+                keep_indices = nms_custom(boxes=boxes, scores=scores, iou_threshold=iou_threshold)
             else:
-                keep_indices = torch.ops.torchvision.nms(boxes=boxes, scores=scores, iou_threshold=iou_threshold)
+                # Use torchvision NMS with correct format
+                keep_indices = torchvision.ops.nms(boxes, scores, iou_threshold)
             
             boxes = boxes[keep_indices]
             scores = scores[keep_indices]
             labels = labels[keep_indices]
         
-        if len(boxes.shape) == 1:
-            boxes = boxes.unsqueeze(0)
-            scores = scores.unsqueeze(0)
-            labels = labels.unsqueeze(0)
-            
+        # Visualize results
         output = visualize_bbox(input_img, boxes, labels, scores, classes_map, alpha=alpha)
-        detection_info = f"Found {len(boxes)} detections after NMS ({nms_method})"
-        return output, detection_info
+        info = f"✅ Found {len(boxes)} detections | NMS: {nms_method} | Threshold: {conf_threshold:.2f}"
+        
+        return output, info
             
     except Exception as e:
-        print(f"[ERROR] recognize_image failed: {e}")
-        error_msg = f"Error during processing: {str(e)}"
+        print(f"[ERROR] process_image failed: {e}")
+        error_msg = f"❌ Processing error: {str(e)}"
         if input_img is not None:
             return np.array(input_img), error_msg
         return np.zeros((512, 512, 3), dtype=np.uint8), error_msg
 
-def gradio_reset():
+def reset_interface():
+    """Reset all interface components."""
     return gr.update(value=None), gr.update(value=None), gr.update(value="")
 
+def create_legend_html():
+    """Create HTML for the class legend."""
+    cmap = colormap(N=len(classes_map), normalized=False)
+    legend_items = []
+    
+    for class_id, class_name in classes_map.items():
+        color_rgb = cmap[class_id % len(cmap)]
+        color_hex = f"#{color_rgb[0]:02x}{color_rgb[1]:02x}{color_rgb[2]:02x}"
+        legend_items.append(f"""
+        <div style='display: flex; align-items: center; padding: 8px; margin: 4px; background-color: #f8f9fa; border-radius: 6px;'>
+            <div style='width: 24px; height: 24px; background-color: {color_hex}; margin-right: 12px; border: 2px solid #dee2e6; border-radius: 4px;'></div>
+            <span style='font-weight: 500; color: #495057;'>{class_name}</span>
+        </div>
+        """)
+    
+    return f"""
+    <div style='display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 8px; padding: 16px; background-color: #ffffff; border-radius: 8px; border: 1px solid #e9ecef;'>
+        {''.join(legend_items)}
+    </div>
+    """
+
 if __name__ == "__main__":
-    print(f"Using device: {device}")
+    print(f"🚀 Starting Document Layout Analysis App")
+    print(f"📱 Device: {device}")
+    print(f"🤖 Available models: {len(MODELS)}")
     
-    # Custom CSS for better scrolling and layout
+    # Custom CSS for full-width layout
     custom_css = """
     .gradio-container {
-        max-width: 1200px !important;
-        margin: auto !important;
+        max-width: 100% !important;
+        padding: 20px !important;
     }
-    .main-content {
-        overflow-y: auto !important;
-        max-height: 100vh !important;
+    
+    .main-container {
+        width: 100% !important;
+        max-width: none !important;
     }
+    
+    .panel-left, .panel-right {
+        min-height: 600px;
+        padding: 20px;
+        background: #f8f9fa;
+        border-radius: 12px;
+        border: 1px solid #e9ecef;
+    }
+    
+    .control-section {
+        margin-bottom: 20px;
+        padding: 15px;
+        background: white;
+        border-radius: 8px;
+        border: 1px solid #dee2e6;
+    }
+    
+    .status-good { color: #28a745; font-weight: bold; }
+    .status-error { color: #dc3545; font-weight: bold; }
+    .status-info { color: #17a2b8; font-weight: bold; }
     """
     
-    with gr.Blocks(title="Document Layout Analysis", theme=gr.themes.Soft(), css=custom_css) as demo:
+    # Create Gradio interface
+    with gr.Blocks(
+        title="📄 Document Layout Analysis - Full Width", 
+        theme=gr.themes.Soft(),
+        css=custom_css
+    ) as demo:
+        
         # Header
         gr.HTML("""
-        <div style="text-align: center; margin-bottom: 20px;">
-            <h1>🔍 Document Layout Analysis</h1>
-            <p>Using Docling Layout Models for document structure detection</p>
+        <div style='text-align: center; padding: 30px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 15px; margin-bottom: 30px;'>
+            <h1 style='margin: 0; font-size: 3em; text-shadow: 2px 2px 4px rgba(0,0,0,0.3);'>🔍 Document Layout Analysis</h1>
+            <p style='margin: 10px 0 0 0; font-size: 1.3em; opacity: 0.9;'>Advanced document structure detection with multiple AI models</p>
         </div>
         """)
         
+        # Main content in two columns
         with gr.Row():
-            # Left Column - Controls
-            with gr.Column(scale=1):
-                # Model selection
-                model_dropdown = gr.Dropdown(
-                    choices=list(MODELS.keys()),
-                    value="Egret XLarge",
-                    label="🤖 Select Model"
-                )
-                
-                load_btn = gr.Button("📥 Load Model", variant="secondary", size="sm")
-                model_status = gr.Textbox(label="Model Status", interactive=False, value="No model loaded", max_lines=2)
+            # LEFT COLUMN - Controls and Input
+            with gr.Column(scale=1, elem_classes=["panel-left"]):
                 
-                input_img = gr.Image(label="📄 Upload Image", type="pil", height=300)
+                # Model Section
+                with gr.Group(elem_classes=["control-section"]):
+                    gr.HTML("<h3>🤖 Model Configuration</h3>")
+                    
+                    model_dropdown = gr.Dropdown(
+                        choices=list(MODELS.keys()),
+                        value="Egret XLarge",
+                        label="Select Model",
+                        info="Choose the AI model for document analysis",
+                        interactive=True
+                    )
+                    
+                    with gr.Row():
+                        load_btn = gr.Button("📥 Load Model", variant="primary", scale=1)
+                        clear_btn = gr.Button("🗑️ Clear All", variant="secondary", scale=1)
+                    
+                    model_status = gr.Textbox(
+                        label="Model Status",
+                        value="��� No model loaded. Please select and load a model.",
+                        interactive=False,
+                        lines=2
+                    )
                 
-                with gr.Row():
-                    clear = gr.Button("🗑️ Clear", size="sm")
-                    predict = gr.Button("🔍 Detect", variant="primary", size="sm")
+                # Image Upload Section
+                with gr.Group(elem_classes=["control-section"]):
+                    gr.HTML("<h3>📄 Image Input</h3>")
+                    
+                    input_img = gr.Image(
+                        label="Upload Document Image",
+                        type="pil",
+                        height=400,
+                        interactive=True
+                    )
+                    
+                    detect_btn = gr.Button("🔍 Analyze Document", variant="primary", size="lg")
                 
-                # Parameters
-                conf_threshold = gr.Slider(0.0, 1.0, value=0.6, step=0.05, label="Confidence Threshold")
-                iou_threshold = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="NMS IoU Threshold")
-                nms_method = gr.Radio(["Custom IoMin", "Standard IoU"], value="Custom IoMin", label="NMS Method")
-                alpha_slider = gr.Slider(0.0, 1.0, value=0.3, step=0.1, label="Overlay Transparency")
+                # Parameters Section  
+                with gr.Group(elem_classes=["control-section"]):
+                    gr.HTML("<h3>⚙️ Detection Parameters</h3>")
+                    
+                    conf_threshold = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.6,
+                        step=0.05,
+                        label="Confidence Threshold",
+                        info="Minimum confidence for detections"
+                    )
+                    
+                    iou_threshold = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.5,
+                        step=0.05,
+                        label="NMS IoU Threshold",
+                        info="Non-maximum suppression threshold"
+                    )
+                    
+                    nms_method = gr.Radio(
+                        choices=["Custom IoMin", "Standard IoU"],
+                        value="Custom IoMin",
+                        label="NMS Algorithm",
+                        info="Choose suppression method"
+                    )
+                    
+                    alpha_slider = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.3,
+                        step=0.1,
+                        label="Overlay Transparency",
+                        info="Transparency of detection overlays"
+                    )
+            
+            # RIGHT COLUMN - Results and Output
+            with gr.Column(scale=1, elem_classes=["panel-right"]):
                 
-            # Right Column - Results
-            with gr.Column(scale=1):
-                gr.HTML("<h3>🎯 Detection Results</h3>")
-                output_img = gr.Image(label="Detected Layout", interactive=False, type="numpy", height=400)
-                detection_info = gr.Textbox(label="Detection Info", interactive=False, max_lines=2)
+                # Results Section
+                with gr.Group(elem_classes=["control-section"]):
+                    gr.HTML("<h3>🎯 Detection Results</h3>")
+                    
+                    output_img = gr.Image(
+                        label="Analyzed Document",
+                        type="numpy",
+                        height=500,
+                        interactive=False
+                    )
+                    
+                    detection_info = gr.Textbox(
+                        label="Analysis Summary",
+                        value="",
+                        interactive=False,
+                        lines=3,
+                        placeholder="Detection results will appear here..."
+                    )
         
-        # Legend at the bottom
-        with gr.Accordion("📋 Detected Classes", open=False):
-            cmap = colormap(N=len(classes_map), normalized=False)
-            legend_items = []
-            for class_id, class_name in classes_map.items():
-                color_rgb = cmap[class_id % len(cmap)]
-                color_hex = f"#{color_rgb[0]:02x}{color_rgb[1]:02x}{color_rgb[2]:02x}"
-                legend_items.append(f'<span style="display:inline-block;width:15px;height:15px;background-color:{color_hex};margin-right:5px;border:1px solid #ccc;"></span>{class_name}')
-            
-            legend_html = f"""
-            <div style='display: grid; grid-template-columns: repeat(3, 1fr); gap: 10px; font-size: 14px;'>
-                {''.join([f'<div>{item}</div>' for item in legend_items])}
-            </div>
-            """
-            gr.HTML(legend_html)
+        # Legend Section (Full Width)
+        with gr.Group():
+            with gr.Accordion("📋 Class Legend - All Detectable Elements", open=False):
+                gr.HTML(create_legend_html())
+        
+        # Event Handlers
+        load_btn.click(
+            fn=load_model,
+            inputs=[model_dropdown],
+            outputs=[model_status]
+        )
+        
+        clear_btn.click(
+            fn=reset_interface,
+            outputs=[input_img, output_img, detection_info]
+        )
         
-        # Event handlers
-        load_btn.click(load_model, inputs=[model_dropdown], outputs=[model_status])
-        clear.click(gradio_reset, inputs=None, outputs=[input_img, output_img, detection_info])
-        predict.click(
-            recognize_image, 
-            inputs=[input_img, conf_threshold, iou_threshold, nms_method, alpha_slider], 
+        detect_btn.click(
+            fn=process_image,
+            inputs=[input_img, conf_threshold, iou_threshold, nms_method, alpha_slider],
             outputs=[output_img, detection_info]
         )
     
-    # Launch
-    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True, share=False)
+    # Launch application
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        debug=True,
+        share=False,
+        show_error=True,
+        inbrowser=True
+    )