Spaces:
Sleeping
Sleeping
| 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): | |
| if isinstance(image_path, Image.Image) or isinstance(image_path, np.ndarray): | |
| image = np.array(image_path) | |
| 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) | |
| 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 | |