Update visualization.py

visualization.py

@@ -43,48 +43,40 @@ def colormap(N=256, normalized=False):
     return cmap
 
 def visualize_bbox(image_path, bboxes, classes, scores, id_to_names, alpha=0.3):
-    """
-    Visualize layout detection results on an image.
-
-    Args:
-        image_path (str): Path to the input image.
-        bboxes (list): List of bounding boxes, each represented as [x_min, y_min, x_max, y_max].
-        classes (list): List of class IDs corresponding to the bounding boxes.
-        id_to_names (dict): Dictionary mapping class IDs to class names.
-        alpha (float): Transparency factor for the filled color (default is 0.3).
-
-    Returns:
-        np.ndarray: Image with visualized layout detection results.
-    """
-    # Check if image_path is a PIL.Image.Image object
     if isinstance(image_path, Image.Image) or isinstance(image_path, np.ndarray):
         image = np.array(image_path)
-        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert RGB to BGR for OpenCV
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
     else:
         image = cv2.imread(image_path)
+        if image is None:
+            raise ValueError(f"Could not load image from path: {image_path}")
 
     overlay = image.copy()
-    
     cmap = colormap(N=len(id_to_names), normalized=False)
-    
-    # Iterate over each bounding box
-    for i, bbox in enumerate(bboxes):
-        x_min, y_min, x_max, y_max = map(int, bbox)
-        class_id = int(classes[i])
-        class_name = id_to_names[class_id]
-        
-        text = class_name + f":{scores[i]:.3f}"
-        
-        color = tuple(int(c) for c in cmap[class_id])
-        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)
-        
-        # Add the class name with a background rectangle
-        (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 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.9, (255, 255, 255), 2)
-    
-    # Blend the overlay with the original image
+
+    if len(bboxes) == 0:
+        print("No bounding boxes to display.")
+        return image  # Return original image if nothing detected
+
+    for i in range(len(bboxes)):
+        try:
+            x_min, y_min, x_max, y_max = map(int, bboxes[i])
+            class_id = int(classes[i])
+            class_name = id_to_names.get(class_id, f"unknown_{class_id}")
+            score = scores[i]
+
+            text = f"{class_name}:{score:.3f}"
+            color = tuple(int(c) for c in cmap[class_id % len(cmap)])
+
+            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)
+
+            (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 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.9, (255, 255, 255), 2)
+
+        except Exception as e:
+            print(f"Skipping box {i} due to error: {e}")
+
     cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
-    
-    return image
+    return image