Update app.py

app.py

@@ -2,42 +2,57 @@ import gradio as gr
 import cv2
 import numpy as np
 import onnxruntime as ort
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, list_repo_files
 
-# --- STEP 1: Download the ONNX Model ---
-print("Downloading ONNX model...")
-model_path = hf_hub_download(repo_id="alex-dinh/PP-DocLayoutV3-ONNX", filename="model.onnx")
+# --- STEP 1: Find the correct ONNX filename ---
+REPO_ID = "alex-dinh/PP-DocLayoutV3-ONNX"
+print(f"Searching for ONNX model in {REPO_ID}...")
+
+# Get list of all files in the repo
+all_files = list_repo_files(repo_id=REPO_ID)
+
+# Find the first file that ends with .onnx
+onnx_filename = next((f for f in all_files if f.endswith('.onnx')), None)
+
+if onnx_filename is None:
+    raise FileNotFoundError(f"No .onnx file found in {REPO_ID}. Repo contents: {all_files}")
+
+print(f"Found model file: {onnx_filename}")
+
+# --- STEP 2: Download the Model ---
+model_path = hf_hub_download(repo_id=REPO_ID, filename=onnx_filename)
 print(f"Model downloaded to: {model_path}")
 
-# --- STEP 2: Initialize ONNX Engine ---
-# This loads the AI "brain" directly without needing Paddle
+# --- STEP 3: Initialize ONNX Engine ---
 session = ort.InferenceSession(model_path)
 input_names = [i.name for i in session.get_inputs()]
 output_names = [o.name for o in session.get_outputs()]
 
-# Define labels map (Standard for PP-DocLayout)
-LABELS = {1: "Text", 2: "Title", 3: "List", 4: "Table", 5: "Figure"}
+# Standard PP-DocLayoutV3 labels
+# Based on the paper/standard V3 config
+LABELS = {0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}
 
-def preprocess_image(image, target_size=(800, 800)):
+def preprocess_image(image, target_size=(640, 640)):
     """
-    Prepares the image exactly how the AI expects it (Resize -> Normalize).
+    Standard RT-DETR / Paddle Preprocessing.
     """
     h, w = image.shape[:2]
     
-    # 1. Resize
-    # We do NOT keep aspect ratio for the input blob, but we keep scales to fix boxes later
-    img_resized = cv2.resize(image, target_size)
+    # 1. Resize (Standard size for V3 is usually 640x640 or 800x800)
+    # We use linear interpolation
+    img_resized = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
     
-    # 2. Normalize (Standard ImageNet mean/std)
+    # 2. Normalize (ImageNet mean/std)
     img_data = img_resized.astype(np.float32) / 255.0
     mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
     std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
     img_data = (img_data - mean) / std
     
-    # 3. Transpose to (Batch, Channel, Height, Width)
+    # 3. Transpose (HWC -> CHW)
     img_data = img_data.transpose(2, 0, 1)[None, :, :, :]
     
-    # Calculate scale factors to map detections back to original image
+    # Scale factor for post-processing
+    # shape: [batch, 2] -> [height_scale, width_scale]
     scale_factor = np.array([target_size[0] / h, target_size[1] / w], dtype=np.float32).reshape(1, 2)
     
     return img_data, scale_factor
@@ -46,60 +61,75 @@ def analyze_layout(input_image):
     if input_image is None:
         return None, "No image uploaded"
 
-    # Convert PIL to Numpy/OpenCV
     image_np = np.array(input_image)
-    orig_h, orig_w = image_np.shape[:2]
 
     # --- INFERENCE ---
+    # Prepare input
     input_blob, scale_factor = preprocess_image(image_np)
     
-    # ONNX Runtime inputs
+    # Run ONNX (No Paddle dependency!)
     inputs = {
-        input_names[0]: input_blob,           # The image data
-        input_names[1]: scale_factor          # The resize scale
+        input_names[0]: input_blob,
+        input_names[1]: scale_factor # Some exports require this, others don't. 
     }
     
-    # Run!
+    # Handle exports that don't use scale_factor input
+    if len(input_names) == 1:
+        del inputs[input_names[1]]
+    
     outputs = session.run(output_names, inputs)
     
-    # --- POST-PROCESSING ---
-    # Output format is typically [Batch, N, 6] -> [Class, Score, X1, Y1, X2, Y2]
-    detections = outputs[0]  
+    # --- PARSE RESULTS ---
+    # Output format varies by export, but usually it's [Batch, N, 6] 
+    # [Class, Score, X1, Y1, X2, Y2]
+    detections = outputs[0] 
     
     viz_image = image_np.copy()
     log = []
 
+    # If the output is wrapped in a batch dimension, unwrap it
+    if len(detections.shape) == 3:
+        detections = detections[0]
+
     for det in detections:
-        class_id = int(det[0])
+        # Check confidence (usually index 1)
         score = det[1]
-        bbox = det[2:]
+        if score < 0.4: continue 
 
-        if score < 0.5: continue # Filter weak detections
+        class_id = int(det[0])
+        bbox = det[2:]
 
         # Map labels
-        label_name = LABELS.get(class_id, "Unknown")
+        label_name = LABELS.get(class_id, f"Class {class_id}")
         
-        # Coordinates
-        x1, y1, x2, y2 = map(int, bbox)
-        
-        # Color coding
-        color = (0, 255, 0) # Green
-        if label_name == "Title": color = (0, 0, 255)
-        elif label_name == "Table": color = (255, 255, 0)
-        elif label_name == "Figure": color = (255, 0, 0)
-
-        # Draw
-        cv2.rectangle(viz_image, (x1, y1), (x2, y2), color, 3)
-        cv2.putText(viz_image, f"{label_name} {score:.2f}", (x1, y1-10), 
-                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
-        
-        log.append(f"Found {label_name} at [{x1}, {y1}, {x2}, {y2}] (Conf: {score:.2f})")
+        # Draw Box
+        try:
+            x1, y1, x2, y2 = map(int, bbox)
+            
+            # Color coding
+            color = (0, 255, 0) # Green
+            if "Title" in label_name: color = (0, 0, 255)
+            elif "Table" in label_name: color = (255, 255, 0)
+            elif "Figure" in label_name: color = (255, 0, 0)
+
+            cv2.rectangle(viz_image, (x1, y1), (x2, y2), color, 3)
+            
+            # Label
+            label_text = f"{label_name} {score:.2f}"
+            (w, h), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+            cv2.rectangle(viz_image, (x1, y1 - 20), (x1 + w, y1), color, -1)
+            cv2.putText(viz_image, label_text, (x1, y1 - 5), 
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+            
+            log.append(f"Found {label_name} at [{x1}, {y1}, {x2}, {y2}]")
+        except Exception:
+            pass
 
     return viz_image, "\n".join(log)
 
 with gr.Blocks(title="ONNX Layout Analysis") as demo:
     gr.Markdown("## ⚡ Fast V3 Layout Analysis (ONNX)")
-    gr.Markdown("Uses **PP-DocLayoutV3** via ONNX Runtime. No Paddle dependencies.")
+    gr.Markdown(f"Auto-loaded model from `{REPO_ID}`")
     
     with gr.Row():
         with gr.Column():