Updated detection module script and requirements.txt for ONNX models than Keras

requirements.txt

@@ -1,9 +1,9 @@
-tensorflow-cpu==2.15.0
-gradio
-fastapi
-uvicorn
+onnxruntime==1.19.0
 numpy
 pillow
-matplotlib
 opencv-python-headless
-python-multipart
+fastapi
+uvicorn
+gradio
+python-multipart
+matplotlib

visualization.py

@@ -1,4 +1,4 @@
-import tensorflow as tf
+import onnxruntime as ort
 import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
@@ -11,76 +11,9 @@ from typing import List, Tuple, Dict, Optional
 
 
 # ============================================================================
-# CUSTOM LOSS CLASS (Required for model loading)
+# CONFIGURATION - UPDATED FOR ONNX
 # ============================================================================
-@tf.keras.utils.register_keras_serializable()
-class LossCalculation(tf.keras.losses.Loss):
-    """Custom loss function for wireframe detection."""
-
-    def __init__(self, num_classes=7, lambda_coord=5.0, lambda_noobj=0.5,
-                 name='loss_calculation', reduction='sum_over_batch_size', **kwargs):
-        super().__init__(name=name, reduction=reduction)
-        self.num_classes = num_classes
-        self.lambda_coord = lambda_coord
-        self.lambda_noobj = lambda_noobj
-
-    def call(self, y_true, y_pred):
-        obj_true = y_true[..., 0]
-        box_true = y_true[..., 1:5]
-        cls_true = y_true[..., 5:]
-
-        obj_pred_logits = y_pred[..., 0]
-        box_pred = y_pred[..., 1:5]
-        cls_pred_logits = y_pred[..., 5:]
-
-        obj_mask = tf.cast(obj_true > 0.5, tf.float32)
-        noobj_mask = 1.0 - obj_mask
-        num_pos = tf.maximum(tf.reduce_sum(obj_mask), 1.0)
-
-        obj_loss_pos = obj_mask * tf.nn.sigmoid_cross_entropy_with_logits(
-            labels=obj_true, logits=obj_pred_logits)
-        obj_loss_neg = noobj_mask * tf.nn.sigmoid_cross_entropy_with_logits(
-            labels=obj_true, logits=obj_pred_logits)
-        obj_loss = (tf.reduce_sum(obj_loss_pos) + self.lambda_noobj * tf.reduce_sum(obj_loss_neg)) / tf.cast(
-            tf.size(obj_true), tf.float32)
-
-        xy_pred = tf.nn.sigmoid(box_pred[..., 0:2])
-        wh_pred = tf.nn.sigmoid(box_pred[..., 2:4])
-        xy_true = box_true[..., 0:2]
-        wh_true = box_true[..., 2:4]
-
-        xy_loss = tf.reduce_sum(obj_mask[..., tf.newaxis] * self._smooth_l1_loss(xy_true - xy_pred)) / num_pos
-        wh_loss = tf.reduce_sum(obj_mask[..., tf.newaxis] * self._smooth_l1_loss(wh_true - wh_pred)) / num_pos
-        box_loss = self.lambda_coord * (xy_loss + wh_loss)
-
-        cls_loss = tf.reduce_sum(obj_mask * tf.nn.softmax_cross_entropy_with_logits(
-            labels=cls_true, logits=cls_pred_logits)) / num_pos
-
-        total_loss = obj_loss + box_loss + cls_loss
-        return tf.clip_by_value(total_loss, 0.0, 100.0)
-
-    def _smooth_l1_loss(self, x, beta=1.0):
-        abs_x = tf.abs(x)
-        return tf.where(abs_x < beta, 0.5 * x * x / beta, abs_x - 0.5 * beta)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update({
-            'num_classes': self.num_classes,
-            'lambda_coord': self.lambda_coord,
-            'lambda_noobj': self.lambda_noobj,
-        })
-        return config
-
-    @classmethod
-    def from_config(cls, config):
-        return cls(**config)
-
-
-# ============================================================================
-# CONFIGURATION - UPDATED FOR BETTER PRECISION
-# ============================================================================
-MODEL_PATH = "./wireframe_detection_model_best_700.keras"
+MODEL_PATH = "./wireframe_detection_model_best_700.onnx"  # Changed to .onnx
 OUTPUT_DIR = "./output/"
 CLASS_NAMES = ["button", "checkbox", "image", "navbar", "paragraph", "text", "textfield"]
 
@@ -88,27 +21,102 @@ IMG_SIZE = 416
 CONF_THRESHOLD = 0.1
 IOU_THRESHOLD = 0.1
 
-# Layout Configuration - INCREASED GRID DENSITY
-GRID_COLUMNS = 24  # Doubled from 12 for finer precision
-ALIGNMENT_THRESHOLD = 10  # Reduced from 15 for tighter alignment
-SIZE_CLUSTERING_THRESHOLD = 15  # Reduced from 20 for better size grouping
+# Layout Configuration
+GRID_COLUMNS = 24
+ALIGNMENT_THRESHOLD = 10
+SIZE_CLUSTERING_THRESHOLD = 15
 
-# Standard sizes for each element type (relative units) - UPDATED FOR SMALLER BUTTONS/CHECKBOXES
+# Standard sizes for each element type (relative units)
 STANDARD_SIZES = {
-    'button': {'width': 2, 'height': 1},  # Smaller button (was 2x1, now in finer grid)
-    'checkbox': {'width': 1, 'height': 1},  # Keep small checkbox
-    'textfield': {'width': 5, 'height': 1},  # Adjusted for new grid
-    'text': {'width': 3, 'height': 1},  # Adjusted
-    'paragraph': {'width': 8, 'height': 2},  # Adjusted
-    'image': {'width': 4, 'height': 4},  # Adjusted
-    'navbar': {'width': 24, 'height': 1}  # Full width in new grid
+    'button': {'width': 2, 'height': 1},
+    'checkbox': {'width': 1, 'height': 1},
+    'textfield': {'width': 5, 'height': 1},
+    'text': {'width': 3, 'height': 1},
+    'paragraph': {'width': 8, 'height': 2},
+    'image': {'width': 4, 'height': 4},
+    'navbar': {'width': 24, 'height': 1}
 }
 
-model = None
+ort_session = None  # Changed from model to ort_session
 
 
 # ============================================================================
-# DATA STRUCTURES
+# UTILITY FUNCTIONS FOR ONNX
+# ============================================================================
+def sigmoid(x):
+    """Sigmoid activation function."""
+    return 1 / (1 + np.exp(-np.clip(x, -500, 500)))
+
+
+def softmax(x, axis=-1):
+    """Softmax activation function."""
+    exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
+    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
+
+
+def non_max_suppression_numpy(boxes, scores, iou_threshold=0.5, score_threshold=0.1):
+    """
+    Pure NumPy implementation of Non-Maximum Suppression.
+    
+    Args:
+        boxes: Array of shape (N, 4) with format [x1, y1, x2, y2]
+        scores: Array of shape (N,) with confidence scores
+        iou_threshold: IoU threshold for suppression
+        score_threshold: Minimum score threshold
+    
+    Returns:
+        List of indices to keep
+    """
+    if len(boxes) == 0:
+        return []
+    
+    # Filter by score threshold
+    keep_mask = scores >= score_threshold
+    boxes = boxes[keep_mask]
+    scores = scores[keep_mask]
+    
+    if len(boxes) == 0:
+        return []
+    
+    # Get coordinates
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    
+    # Calculate areas
+    areas = (x2 - x1) * (y2 - y1)
+    
+    # Sort by scores
+    order = scores.argsort()[::-1]
+    
+    keep = []
+    while order.size > 0:
+        # Pick the box with highest score
+        i = order[0]
+        keep.append(i)
+        
+        # Calculate IoU with remaining boxes
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+        
+        w = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        
+        intersection = w * h
+        iou = intersection / (areas[i] + areas[order[1:]] - intersection)
+        
+        # Keep boxes with IoU less than threshold
+        inds = np.where(iou <= iou_threshold)[0]
+        order = order[inds + 1]
+    
+    return keep
+
+
+# ============================================================================
+# DATA STRUCTURES (unchanged)
 # ============================================================================
 @dataclass
 class Element:
@@ -139,13 +147,13 @@ class NormalizedElement:
 
 
 # ============================================================================
-# PREDICTION EXTRACTION
+# PREDICTION EXTRACTION - MODIFIED FOR ONNX
 # ============================================================================
 def get_predictions(image_path: str) -> Tuple[Image.Image, List[Element]]:
-    """Extract predictions from the model."""
-    global model
-    if model is None:
-        raise ValueError("Model not loaded. Please load the model first.")
+    """Extract predictions from the ONNX model."""
+    global ort_session
+    if ort_session is None:
+        raise ValueError("ONNX model not loaded. Please load the model first.")
 
     # Load and preprocess image
     pil_img = Image.open(image_path).convert("RGB")
@@ -155,25 +163,28 @@ def get_predictions(image_path: str) -> Tuple[Image.Image, List[Element]]:
     img_array = np.array(resized_img, dtype=np.float32) / 255.0
     input_tensor = np.expand_dims(img_array, axis=0)
 
-    # Get predictions
-    pred_grid = model.predict(input_tensor, verbose=0)[0]
+    # Get predictions from ONNX model
+    input_name = ort_session.get_inputs()[0].name
+    output_name = ort_session.get_outputs()[0].name
+    pred_grid = ort_session.run([output_name], {input_name: input_tensor})[0][0]
+    
     raw_boxes = []
     S = pred_grid.shape[0]
     cell_size = 1.0 / S
 
     for row in range(S):
         for col in range(S):
-            obj_score = float(tf.nn.sigmoid(pred_grid[row, col, 0]))
+            obj_score = float(sigmoid(pred_grid[row, col, 0]))
             if obj_score < CONF_THRESHOLD:
                 continue
 
-            x_offset = float(tf.nn.sigmoid(pred_grid[row, col, 1]))
-            y_offset = float(tf.nn.sigmoid(pred_grid[row, col, 2]))
-            width = float(tf.nn.sigmoid(pred_grid[row, col, 3]))
-            height = float(tf.nn.sigmoid(pred_grid[row, col, 4]))
+            x_offset = float(sigmoid(pred_grid[row, col, 1]))
+            y_offset = float(sigmoid(pred_grid[row, col, 2]))
+            width = float(sigmoid(pred_grid[row, col, 3]))
+            height = float(sigmoid(pred_grid[row, col, 4]))
 
             class_logits = pred_grid[row, col, 5:]
-            class_probs = tf.nn.softmax(class_logits).numpy()
+            class_probs = softmax(class_logits)
             class_id = int(np.argmax(class_probs))
             class_conf = float(class_probs[class_id])
             final_score = obj_score * class_conf
@@ -191,25 +202,24 @@ def get_predictions(image_path: str) -> Tuple[Image.Image, List[Element]]:
             if x2 > x1 and y2 > y1:
                 raw_boxes.append((class_id, final_score, x1, y1, x2, y2))
 
-    # Apply NMS per class
+    # Apply NMS per class using NumPy implementation
     elements = []
     for class_id in range(len(CLASS_NAMES)):
         class_boxes = [(score, x1, y1, x2, y2) for cid, score, x1, y1, x2, y2 in raw_boxes if cid == class_id]
         if not class_boxes:
             continue
 
-        scores = [b[0] for b in class_boxes]
-        boxes_xyxy = [[b[1], b[2], b[3], b[4]] for b in class_boxes]
+        scores = np.array([b[0] for b in class_boxes])
+        boxes_xyxy = np.array([[b[1], b[2], b[3], b[4]] for b in class_boxes])
 
-        selected_indices = tf.image.non_max_suppression(
+        selected_indices = non_max_suppression_numpy(
             boxes=boxes_xyxy,
             scores=scores,
-            max_output_size=50,
             iou_threshold=IOU_THRESHOLD,
             score_threshold=CONF_THRESHOLD
         )
 
-        for idx in selected_indices.numpy():
+        for idx in selected_indices:
             score, x1, y1, x2, y2 = class_boxes[idx]
             elements.append(Element(
                 label=CLASS_NAMES[class_id],
@@ -221,7 +231,7 @@ def get_predictions(image_path: str) -> Tuple[Image.Image, List[Element]]:
 
 
 # ============================================================================
-# ALIGNMENT DETECTION
+# ALIGNMENT DETECTION (unchanged)
 # ============================================================================
 class AlignmentDetector:
     """Detects alignment relationships between elements."""
@@ -333,7 +343,7 @@ class AlignmentDetector:
 
 
 # ============================================================================
-# SIZE NORMALIZATION - UPDATED TO RESPECT ACTUAL SIZES MORE
+# SIZE NORMALIZATION (unchanged)
 # ============================================================================
 class SizeNormalizer:
     """Normalizes element sizes based on type and clustering."""
@@ -384,17 +394,13 @@ class SizeNormalizer:
         return clusters
 
     def get_normalized_size(self, element: Element, size_cluster: List[Element]) -> Tuple[float, float]:
-        """Get normalized size for an element based on its cluster - PRESERVES ACTUAL SIZE BETTER."""
-        # Use the actual detected size instead of aggressive averaging
-        # Only normalize if there's a significant cluster
+        """Get normalized size for an element based on its cluster."""
         if len(size_cluster) >= 3:
-            # Use median instead of mean to avoid outliers
             widths = sorted([e.width for e in size_cluster])
             heights = sorted([e.height for e in size_cluster])
             median_width = widths[len(widths) // 2]
             median_height = heights[len(heights) // 2]
 
-            # Only normalize if element is within 30% of median
             if abs(element.width - median_width) / median_width < 0.3:
                 normalized_width = round(median_width)
             else:
@@ -405,7 +411,6 @@ class SizeNormalizer:
             else:
                 normalized_height = round(element.height)
         else:
-            # Small cluster - keep original size
             normalized_width = round(element.width)
             normalized_height = round(element.height)
 
@@ -413,7 +418,7 @@ class SizeNormalizer:
 
 
 # ============================================================================
-# GRID-BASED LAYOUT SYSTEM - UPDATED FOR FINER PRECISION
+# GRID-BASED LAYOUT SYSTEM (unchanged)
 # ============================================================================
 class GridLayoutSystem:
     """Grid-based layout system for precise positioning."""
@@ -432,44 +437,36 @@ class GridLayoutSystem:
         print(f"📐 Cell size: {self.cell_width:.1f}px × {self.cell_height:.1f}px")
 
     def snap_to_grid(self, bbox: List[float], element_label: str, preserve_size: bool = True) -> List[float]:
-        """Snap bounding box to grid - UPDATED TO PRESERVE ORIGINAL SIZE BETTER."""
+        """Snap bounding box to grid."""
         x1, y1, x2, y2 = bbox
         original_width = x2 - x1
         original_height = y2 - y1
 
-        # Calculate center
         center_x = (x1 + x2) / 2
         center_y = (y1 + y2) / 2
 
-        # Find nearest grid cell for center
         center_col = round(center_x / self.cell_width)
         center_row = round(center_y / self.cell_height)
 
         if preserve_size:
-            # Calculate span based on actual size (don't force to standard)
             width_cells = max(1, round(original_width / self.cell_width))
             height_cells = max(1, round(original_height / self.cell_height))
         else:
-            # Use standard size
             standard = STANDARD_SIZES.get(element_label, {'width': 2, 'height': 1})
             width_cells = max(1, round(original_width / self.cell_width))
             height_cells = max(1, round(original_height / self.cell_height))
 
-            # Only adjust to standard if very close
             if abs(width_cells - standard['width']) <= 0.5:
                 width_cells = standard['width']
             if abs(height_cells - standard['height']) <= 0.5:
                 height_cells = standard['height']
 
-        # Calculate start position (center the element)
         start_col = center_col - width_cells // 2
         start_row = center_row - height_cells // 2
 
-        # Clamp to grid bounds
         start_col = max(0, min(start_col, self.num_columns - width_cells))
         start_row = max(0, min(start_row, self.num_rows - height_cells))
 
-        # Convert back to pixels
         snapped_x1 = start_col * self.cell_width
         snapped_y1 = start_row * self.cell_height
         snapped_x2 = (start_col + width_cells) * self.cell_width
@@ -497,7 +494,7 @@ class GridLayoutSystem:
 
 
 # ============================================================================
-# OVERLAP DETECTION & RESOLUTION - UPDATED WITH BETTER STRATEGIES
+# OVERLAP DETECTION & RESOLUTION (unchanged)
 # ============================================================================
 class OverlapResolver:
     """Detects and resolves overlapping elements."""
@@ -506,7 +503,7 @@ class OverlapResolver:
         self.elements = elements
         self.img_width = img_width
         self.img_height = img_height
-        self.overlap_threshold = 0.2  # Reduced from 0.3 - be more aggressive
+        self.overlap_threshold = 0.2
 
     def compute_iou(self, bbox1: List[float], bbox2: List[float]) -> float:
         """Compute Intersection over Union between two bounding boxes."""
@@ -545,7 +542,7 @@ class OverlapResolver:
         return overlap_ratio1, overlap_ratio2
 
     def resolve_overlaps(self, normalized_elements: List[NormalizedElement]) -> List[NormalizedElement]:
-        """Resolve overlaps by adjusting element positions - IMPROVED ALGORITHM."""
+        """Resolve overlaps by adjusting element positions."""
         print("\n🔍 Checking for overlaps...")
 
         overlaps = []
@@ -579,7 +576,6 @@ class OverlapResolver:
 
         print(f"⚠️  Found {len(overlaps)} overlapping element pairs")
 
-        # Sort by overlap severity
         overlaps.sort(key=lambda x: x['overlap'], reverse=True)
 
         elements_to_remove = set()
@@ -595,7 +591,6 @@ class OverlapResolver:
             elem2 = overlap_info['elem2']
             overlap_ratio = overlap_info['overlap']
 
-            # Strategy 1: Nearly complete overlap (>70%) - remove lower confidence
             if overlap_ratio > 0.7:
                 if elem1.original.score < elem2.original.score:
                     elements_to_remove.add(idx1)
@@ -606,13 +601,11 @@ class OverlapResolver:
                     print(f"  🗑️  Removing {elem2.original.label} (conf: {elem2.original.score:.2f}) - "
                           f"overlaps {overlap_ratio * 100:.1f}% with {elem1.original.label}")
 
-            # Strategy 2: Significant overlap (40-70%) - try to separate
             elif overlap_ratio > 0.4:
                 self._try_separate_elements(elem1, elem2, overlap_info)
                 print(f"  ↔️  Separating {elem1.original.label} and {elem2.original.label} "
                       f"(overlap: {overlap_ratio * 100:.1f}%)")
 
-            # Strategy 3: Moderate overlap (20-40%) - shrink slightly
             else:
                 self._shrink_overlapping_edges(elem1, elem2, overlap_info)
                 print(f"  📏 Shrinking {elem1.original.label} and {elem2.original.label} "
@@ -629,11 +622,10 @@ class OverlapResolver:
 
     def _try_separate_elements(self, elem1: NormalizedElement, elem2: NormalizedElement,
                                overlap_info: Dict):
-        """Try to separate two significantly overlapping elements - IMPROVED."""
+        """Try to separate two significantly overlapping elements."""
         bbox1 = elem1.normalized_bbox
         bbox2 = elem2.normalized_bbox
 
-        # Calculate overlap dimensions
         overlap_x1 = max(bbox1[0], bbox2[0])
         overlap_y1 = max(bbox1[1], bbox2[1])
         overlap_x2 = min(bbox1[2], bbox2[2])
@@ -642,45 +634,35 @@ class OverlapResolver:
         overlap_width = overlap_x2 - overlap_x1
         overlap_height = overlap_y2 - overlap_y1
 
-        # Calculate centers
         center1_x = (bbox1[0] + bbox1[2]) / 2
         center1_y = (bbox1[1] + bbox1[3]) / 2
         center2_x = (bbox2[0] + bbox2[2]) / 2
         center2_y = (bbox2[1] + bbox2[3]) / 2
 
-        # Determine separation direction
         dx = abs(center2_x - center1_x)
         dy = abs(center2_y - center1_y)
 
-        # Add minimum gap
-        min_gap = 3  # pixels
+        min_gap = 3
 
         if dx > dy:
-            # Separate horizontally
             if center1_x < center2_x:
-                # elem1 is left of elem2
                 midpoint = (bbox1[2] + bbox2[0]) / 2
                 bbox1[2] = midpoint - min_gap
                 bbox2[0] = midpoint + min_gap
             else:
-                # elem2 is left of elem1
                 midpoint = (bbox2[2] + bbox1[0]) / 2
                 bbox2[2] = midpoint - min_gap
                 bbox1[0] = midpoint + min_gap
         else:
-            # Separate vertically
             if center1_y < center2_y:
-                # elem1 is above elem2
                 midpoint = (bbox1[3] + bbox2[1]) / 2
                 bbox1[3] = midpoint - min_gap
                 bbox2[1] = midpoint + min_gap
             else:
-                # elem2 is above elem1
                 midpoint = (bbox2[3] + bbox1[1]) / 2
                 bbox2[3] = midpoint - min_gap
                 bbox1[1] = midpoint + min_gap
 
-        # Ensure boxes remain valid
         self._ensure_valid_bbox(bbox1)
         self._ensure_valid_bbox(bbox2)
 
@@ -690,7 +672,6 @@ class OverlapResolver:
         bbox1 = elem1.normalized_bbox
         bbox2 = elem2.normalized_bbox
 
-        # Calculate overlap region
         overlap_x1 = max(bbox1[0], bbox2[0])
         overlap_y1 = max(bbox1[1], bbox2[1])
         overlap_x2 = min(bbox1[2], bbox2[2])
@@ -699,11 +680,9 @@ class OverlapResolver:
         overlap_width = overlap_x2 - overlap_x1
         overlap_height = overlap_y2 - overlap_y1
 
-        # Shrink by 50% of overlap plus small gap
-        gap = 2  # pixels
+        gap = 2
 
         if overlap_width > overlap_height:
-            # Horizontal overlap is larger
             shrink = overlap_width / 2 + gap
             if bbox1[0] < bbox2[0]:
                 bbox1[2] -= shrink
@@ -712,7 +691,6 @@ class OverlapResolver:
                 bbox2[2] -= shrink
                 bbox1[0] += shrink
         else:
-            # Vertical overlap is larger
             shrink = overlap_height / 2 + gap
             if bbox1[1] < bbox2[1]:
                 bbox1[3] -= shrink
@@ -726,9 +704,8 @@ class OverlapResolver:
 
     def _ensure_valid_bbox(self, bbox: List[float]):
         """Ensure bounding box has minimum size and is within image bounds."""
-        min_size = 8  # Reduced minimum size
+        min_size = 8
 
-        # Ensure minimum size
         if bbox[2] - bbox[0] < min_size:
             center_x = (bbox[0] + bbox[2]) / 2
             bbox[0] = center_x - min_size / 2
@@ -739,7 +716,6 @@ class OverlapResolver:
             bbox[1] = center_y - min_size / 2
             bbox[3] = center_y + min_size / 2
 
-        # Clamp to image bounds
         bbox[0] = max(0, min(bbox[0], self.img_width))
         bbox[1] = max(0, min(bbox[1], self.img_height))
         bbox[2] = max(0, min(bbox[2], self.img_width))
@@ -747,7 +723,7 @@ class OverlapResolver:
 
 
 # ============================================================================
-# MAIN NORMALIZATION ENGINE
+# MAIN NORMALIZATION ENGINE (unchanged)
 # ============================================================================
 class LayoutNormalizer:
     """Main engine for normalizing wireframe layout."""
@@ -764,7 +740,6 @@ class LayoutNormalizer:
         """Normalize all elements with proper sizing and alignment."""
         print("\n🔧 Starting layout normalization...")
 
-        # Step 1: Detect alignments
         h_alignments = self.alignment_detector.detect_horizontal_alignments()
         v_alignments = self.alignment_detector.detect_vertical_alignments()
         edge_alignments = self.alignment_detector.detect_edge_alignments()
@@ -772,11 +747,9 @@ class LayoutNormalizer:
         print(f"✓ Found {len(h_alignments)} horizontal alignment groups")
         print(f"✓ Found {len(v_alignments)} vertical alignment groups")
 
-        # Step 2: Cluster sizes by type
         size_clusters = self.size_normalizer.cluster_sizes_by_type()
         print(f"✓ Created size clusters for {len(size_clusters)} element types")
 
-        # Step 3: Create element-to-cluster mapping
         element_to_cluster = {}
         element_to_size_category = {}
         for label, clusters in size_clusters.items():
@@ -786,18 +759,14 @@ class LayoutNormalizer:
                     element_to_cluster[id(elem)] = cluster
                     element_to_size_category[id(elem)] = category
 
-        # Step 4: Normalize each element
         normalized_elements = []
 
         for elem in self.elements:
-            # Get size cluster
             cluster = element_to_cluster.get(id(elem), [elem])
             size_category = element_to_size_category.get(id(elem), f"{elem.label}_default")
 
-            # Get normalized size
             norm_width, norm_height = self.size_normalizer.get_normalized_size(elem, cluster)
 
-            # Create normalized bbox (centered on original)
             center_x, center_y = elem.center_x, elem.center_y
             norm_bbox = [
                 center_x - norm_width / 2,
@@ -806,7 +775,6 @@ class LayoutNormalizer:
                 center_y + norm_height / 2
             ]
 
-            # Snap to grid - preserve original size better
             snapped_bbox = self.grid.snap_to_grid(norm_bbox, elem.label, preserve_size=True)
             grid_position = self.grid.get_grid_position(snapped_bbox)
 
@@ -817,12 +785,10 @@ class LayoutNormalizer:
                 size_category=size_category
             ))
 
-        # Step 5: Apply alignment corrections
         normalized_elements = self._apply_alignment_corrections(
             normalized_elements, h_alignments, v_alignments, edge_alignments
         )
 
-        # Step 6: Resolve overlaps
         overlap_resolver = OverlapResolver(self.elements, self.img_width, self.img_height)
         normalized_elements = overlap_resolver.resolve_overlaps(normalized_elements)
 
@@ -835,32 +801,26 @@ class LayoutNormalizer:
                                      edge_alignments: Dict) -> List[NormalizedElement]:
         """Apply alignment corrections to normalized elements."""
 
-        # Create lookup dictionary
         elem_to_normalized = {id(ne.original): ne for ne in normalized_elements}
 
-        # Align horizontally grouped elements
         for h_group in h_alignments:
             norm_group = [elem_to_normalized[id(e)] for e in h_group if id(e) in elem_to_normalized]
             if len(norm_group) > 1:
-                # Align to average Y position
                 avg_y = sum((ne.normalized_bbox[1] + ne.normalized_bbox[3]) / 2 for ne in norm_group) / len(norm_group)
                 for ne in norm_group:
                     height = ne.normalized_bbox[3] - ne.normalized_bbox[1]
                     ne.normalized_bbox[1] = avg_y - height / 2
                     ne.normalized_bbox[3] = avg_y + height / 2
 
-        # Align vertically grouped elements
         for v_group in v_alignments:
             norm_group = [elem_to_normalized[id(e)] for e in v_group if id(e) in elem_to_normalized]
             if len(norm_group) > 1:
-                # Align to average X position
                 avg_x = sum((ne.normalized_bbox[0] + ne.normalized_bbox[2]) / 2 for ne in norm_group) / len(norm_group)
                 for ne in norm_group:
                     width = ne.normalized_bbox[2] - ne.normalized_bbox[0]
                     ne.normalized_bbox[0] = avg_x - width / 2
                     ne.normalized_bbox[2] = avg_x + width / 2
 
-        # Align edges
         for edge_type, groups in edge_alignments.items():
             for edge_group in groups:
                 norm_group = [elem_to_normalized[id(e)] for e in edge_group if id(e) in elem_to_normalized]
@@ -894,7 +854,7 @@ class LayoutNormalizer:
 
 
 # ============================================================================
-# VISUALIZATION & EXPORT
+# VISUALIZATION & EXPORT (unchanged)
 # ============================================================================
 def visualize_comparison(pil_img: Image.Image, elements: List[Element],
                          normalized_elements: List[NormalizedElement],
@@ -903,7 +863,6 @@ def visualize_comparison(pil_img: Image.Image, elements: List[Element],
 
     fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 12))
 
-    # Original detections
     ax1.imshow(pil_img)
     ax1.set_title("Original Predictions", fontsize=16, weight='bold')
     ax1.axis('off')
@@ -918,12 +877,10 @@ def visualize_comparison(pil_img: Image.Image, elements: List[Element],
         ax1.text(x1, y1 - 5, elem.label, color='red', fontsize=8,
                  bbox=dict(facecolor='white', alpha=0.7))
 
-    # Normalized layout
     ax2.imshow(pil_img)
     ax2.set_title("Normalized & Aligned Layout", fontsize=16, weight='bold')
     ax2.axis('off')
 
-    # Draw grid
     for x in range(grid_system.num_columns + 1):
         x_pos = x * grid_system.cell_width
         ax2.axvline(x=x_pos, color='blue', linestyle=':', linewidth=0.5, alpha=0.3)
@@ -931,7 +888,6 @@ def visualize_comparison(pil_img: Image.Image, elements: List[Element],
         y_pos = y * grid_system.cell_height
         ax2.axhline(y=y_pos, color='blue', linestyle=':', linewidth=0.5, alpha=0.3)
 
-    # Draw normalized elements
     np.random.seed(42)
     colors = plt.cm.Set3(np.linspace(0, 1, len(CLASS_NAMES)))
     color_map = {name: colors[i] for i, name in enumerate(CLASS_NAMES)}
@@ -940,14 +896,12 @@ def visualize_comparison(pil_img: Image.Image, elements: List[Element],
         x1, y1, x2, y2 = norm_elem.normalized_bbox
         color = color_map[norm_elem.original.label]
 
-        # Normalized box (thick)
         rect = patches.Rectangle(
             (x1, y1), x2 - x1, y2 - y1,
             linewidth=3, edgecolor=color, facecolor='none'
         )
         ax2.add_patch(rect)
 
-        # Original box (thin, dashed)
         ox1, oy1, ox2, oy2 = norm_elem.original.bbox
         orig_rect = patches.Rectangle(
             (ox1, oy1), ox2 - ox1, oy2 - oy1,
@@ -956,7 +910,6 @@ def visualize_comparison(pil_img: Image.Image, elements: List[Element],
         )
         ax2.add_patch(orig_rect)
 
-        # Label
         grid_pos = norm_elem.grid_position
         label_text = f"{norm_elem.original.label}\n{norm_elem.size_category}\nR{grid_pos['start_row']} C{grid_pos['start_col']}"
         ax2.text(x1 + 5, y1 + 15, label_text, color='white', fontsize=7,
@@ -1086,7 +1039,6 @@ def export_to_html(normalized_elements: List[NormalizedElement],
             text-transform: uppercase;
         }}
 
-        /* Element type specific styles */
         .button {{
             background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
             color: white;
@@ -1160,7 +1112,7 @@ def export_to_html(normalized_elements: List[NormalizedElement],
 </head>
 <body>
     <div class="info-panel">
-        <h3>📏 Layout Info</h3>
+        <h3>📐 Layout Info</h3>
         <p><strong>Grid:</strong> {grid_cols} × {grid_rows}</p>
         <p><strong>Elements:</strong> {total_elements}</p>
         <p><strong>Dimensions:</strong> {img_width}px × {img_height}px</p>
@@ -1208,7 +1160,7 @@ def export_to_html(normalized_elements: List[NormalizedElement],
 
 
 # ============================================================================
-# MAIN PIPELINE
+# MAIN PIPELINE - MODIFIED FOR ONNX
 # ============================================================================
 def process_wireframe(image_path: str,
                       save_json: bool = True,
@@ -1230,51 +1182,46 @@ def process_wireframe(image_path: str,
     print("=== PROCESS_WIREFRAME START ===")
     print("Input image path:", image_path)
     print("File exists:", os.path.exists(image_path))
-    print("File size:", os.path.getsize(image_path))
+    if os.path.exists(image_path):
+        print("File size:", os.path.getsize(image_path))
 
     print("=" * 80)
-    print("🚀 WIREFRAME LAYOUT NORMALIZER")
+    print("🚀 WIREFRAME LAYOUT NORMALIZER (ONNX)")
     print("=" * 80)
 
-    # Step 1: Load model and get predictions
-    global model
-    print("Model object is None?", model is None)
-    print("Model path exists?", os.path.exists(MODEL_PATH))
-    if model is None:
-        print("\n📦 Loading model...")
-        print("Attempting to load keras model:", MODEL_PATH)
-        print("Loaded model summary:")
-        model.summary(print_fn=lambda x: print(x))
+    # Step 1: Load ONNX model and get predictions
+    global ort_session
+    if ort_session is None:
+        print("\n📦 Loading ONNX model...")
+        print("Model path:", MODEL_PATH)
+        print("Model path exists?", os.path.exists(MODEL_PATH))
         try:
-            model = tf.keras.models.load_model(
-                MODEL_PATH,
-                custom_objects={'LossCalculation': LossCalculation}
-            )
-            print("✅ Model loaded successfully!")
+            ort_session = ort.InferenceSession(MODEL_PATH)
+            print("✅ ONNX model loaded successfully!")
+            print(f"Input name: {ort_session.get_inputs()[0].name}")
+            print(f"Input shape: {ort_session.get_inputs()[0].shape}")
+            print(f"Output name: {ort_session.get_outputs()[0].name}")
+            print(f"Output shape: {ort_session.get_outputs()[0].shape}")
         except Exception as e:
-            print(f"❌ Error loading model: {e}")
-            print("\nTrying alternative loading method...")
-            try:
-                model = tf.keras.models.load_model(MODEL_PATH, compile=False)
-                print("✅ Model loaded successfully (without compilation)!")
-            except Exception as e2:
-                print(f"❌ Failed to load model: {e2}")
-                return {}
+            print(f"❌ Error loading ONNX model: {e}")
+            return {}
 
     print(f"\n📸 Processing image: {image_path}")
     print("Running detection inference…")
-    print("Elements detected:", len(elements))
-    for elem in elements:
-        print(" -", elem.label, elem.score, elem.bbox)
-    pil_img, elements = get_predictions(image_path)
-    print(f"✅ Detected {len(elements)} elements")
+    try:
+        pil_img, elements = get_predictions(image_path)
+        print(f"✅ Detected {len(elements)} elements")
+        for elem in elements:
+            print(f"  - {elem.label} (conf: {elem.score:.3f}) at {elem.bbox}")
+    except Exception as e:
+        print(f"❌ Error during prediction: {e}")
+        return {}
 
     if not elements:
-        print("⚠️ No detection output returned.")
-        print("→ Meaning model.predict returned zero raw boxes.")
+        print("⚠️ No elements detected.")
         print("→ Check thresholds:")
-        print("CONF_THRESHOLD:", CONF_THRESHOLD)
-        print("IOU_THRESHOLD:", IOU_THRESHOLD)
+        print(f"   CONF_THRESHOLD: {CONF_THRESHOLD}")
+        print(f"   IOU_THRESHOLD: {IOU_THRESHOLD}")
         return {}
 
     # Step 2: Normalize layout
@@ -1314,7 +1261,6 @@ def process_wireframe(image_path: str,
     print("📊 PROCESSING SUMMARY")
     print("=" * 80)
 
-    # Count by type
     type_counts = {}
     for elem in elements:
         type_counts[elem.label] = type_counts.get(elem.label, 0) + 1
@@ -1323,18 +1269,16 @@ def process_wireframe(image_path: str,
     for elem_type, count in sorted(type_counts.items()):
         print(f"   • {elem_type}: {count}")
 
-    # Size categories
     size_categories = {}
     for norm_elem in normalized_elements:
         size_categories[norm_elem.size_category] = size_categories.get(norm_elem.size_category, 0) + 1
 
-    print(f"\n📐 Size Categories: {len(size_categories)}")
+    print(f"\n📏 Size Categories: {len(size_categories)}")
 
-    # Alignment info
     h_alignments = normalizer.alignment_detector.detect_horizontal_alignments()
     v_alignments = normalizer.alignment_detector.detect_vertical_alignments()
 
-    print(f"\n📍 Alignment:")
+    print(f"\n📐 Alignment:")
     print(f"   • Horizontal groups: {len(h_alignments)}")
     print(f"   • Vertical groups: {len(v_alignments)}")
 
@@ -1383,7 +1327,6 @@ def batch_process(image_dir: str, pattern: str = "*.png"):
                 'error': str(e)
             })
 
-    # Summary
     successful = sum(1 for r in all_results if r['success'])
     print(f"\n{'=' * 80}")
     print(f"📊 BATCH PROCESSING COMPLETE")

wireframe_detection_model_best_700.keras → wireframe.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5dbcc2a9f3222325ee66087ea94a3ac5b6c674844b5173aae30a9f4bf4290f63
-size 53257515
+oid sha256:2c47e1f0f63b4a29dd146331c582860e5981ea0546119b79511a167e856a6277
+size 17701338