Create visualization.py

visualization.py

@@ -0,0 +1,85 @@
+# visualization.py
+import numpy as np
+import cv2
+from PIL import Image
+
+def colormap(N=256, normalized=False):
+    """
+    Generate the color map.
+    Args:
+        N (int): Number of labels (default is 256).
+        normalized (bool): If True, return colors normalized to [0, 1]. Otherwise, return [0, 255].
+    Returns:
+        np.ndarray: Color map array of shape (N, 3).
+    """
+    def bitget(byteval, idx):
+        """
+        Get the bit value at the specified index.
+        Args:
+            byteval (int): The byte value.
+            idx (int): The index of the bit.
+        Returns:
+            int: The bit value (0 or 1).
+        """
+        return ((byteval & (1 << idx)) != 0)
+    
+    cmap = np.zeros((N, 3), dtype=np.uint8)
+    for i in range(N):
+        r = g = b = 0
+        c = i
+        for j in range(8):
+            r = r | (bitget(c, 0) << (7 - j))
+            g = g | (bitget(c, 1) << (7 - j))
+            b = b | (bitget(c, 2) << (7 - j))
+            c = c >> 3
+        cmap[i] = np.array([r, g, b])
+    
+    if normalized:
+        cmap = cmap.astype(np.float32) / 255.0
+    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 or np.ndarray or PIL.Image): Input image or path.
+        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.
+        scores (list): List of confidence scores 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 or numpy array
+    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
+    else:
+        image = cv2.imread(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
+    cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
+    
+    # Convert back to RGB for display
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    return image