Spaces:
Running
Running
| from typing import List | |
| from spatstatapp.paths import find_label_path | |
| from spatstatapp.tile_training_data import load_bboxes | |
| import cv2 | |
| import numpy as np | |
| import supervision as sv | |
| from ultralytics.utils.ops import xywhn2xyxy | |
| def coco_to_detections(label_path, image, id_to_label = {0:"scallop"}): | |
| bboxes = load_bboxes(label_path) | |
| class_id = np.array(bboxes)[:,0].astype(int) | |
| class_name = np.vectorize(id_to_label.get)(class_id) | |
| bboxes = np.array(bboxes)[:,1:] # remove class_id | |
| xyxy = xywhn2xyxy(bboxes, | |
| w = image.shape[1], | |
| h = image.shape[0]) | |
| xyxy = np.clip(xyxy, a_min=0, a_max=max(image.shape)) | |
| detections = sv.Detections( | |
| xyxy=xyxy, | |
| # class_id=np.array(len(xyxy)*[0]), | |
| class_id=class_id, | |
| data = {"class_name":class_name} | |
| ) | |
| return detections | |
| def draw_img_bboxes(img_path): | |
| image = cv2.imread(img_path) | |
| label_path = find_label_path(img_path) | |
| detections = coco_to_detections(label_path, image) | |
| box_annotator = sv.BoxAnnotator() | |
| annotated_frame = box_annotator.annotate( | |
| scene=image.copy(), | |
| detections=detections | |
| ) | |
| return annotated_frame | |
| def stack_images_grid(images): | |
| num_images = len(images) | |
| grid_size = int(np.ceil(np.sqrt(num_images))) | |
| # Calculate the size of the grid | |
| img_height, img_width, _ = images[0].shape | |
| buffer = 2 | |
| grid_height = grid_size * img_height + (grid_size - 1) * buffer | |
| grid_width = grid_size * img_width + (grid_size - 1) * buffer | |
| # Create a white canvas | |
| grid_image = np.ones((grid_height, grid_width, 3), dtype=np.uint8) * 255 | |
| for idx, img in enumerate(images): | |
| row = idx // grid_size | |
| col = idx % grid_size | |
| y = row * (img_height + buffer) | |
| x = col * (img_width + buffer) | |
| grid_image[y:y+img_height, x:x+img_width, :] = img | |
| return grid_image | |
| def show_img_bboxes(img_path): | |
| annotated_frame = draw_img_bboxes(img_path) | |
| sv.plot_image(annotated_frame) | |
| def annotate_DetectionDataset(dataset: sv.DetectionDataset, idxs:List[int]|int): | |
| if isinstance(idxs, int): | |
| idxs = np.random.choice(len(dataset), size=idxs, replace=False).tolist() | |
| frames = [] | |
| for i in idxs: | |
| path, img, detections = dataset[i] | |
| box_annotator = sv.BoxAnnotator() | |
| frames.append(box_annotator.annotate( | |
| scene=img.copy(), | |
| detections=detections | |
| )) | |
| return frames | |
| def show_img_bboxes_grid(img_paths): | |
| """ | |
| shows a grid of images with bounding boxes from list of paths | |
| Args: | |
| img_paths: list of image paths | |
| Example: | |
| show_img_bboxes_grid((data_dir/"images/train").glob("*.jpg")[0:16]) | |
| """ | |
| images = [draw_img_bboxes(img_path) for img_path in img_paths] | |
| annotated_grid = stack_images_grid(images) | |
| sv.plot_image(annotated_grid) | |
| def show_dataset_bboxes_grid(dataset: sv.DetectionDataset, idxs:List[int]|int): | |
| """ | |
| shows a grid of images with bounding boxes from the dataset | |
| Args: | |
| dataset: sv.DetectionDataset: _description_ | |
| idxs: list of indices to show or int: number of random indices to show | |
| Example: | |
| data_dir = Path('dataset_clipped') | |
| dd = sv.DetectionDataset.from_yolo( | |
| images_directory_path= data_dir/"images/train", | |
| annotations_directory_path= data_dir/"labels/train", | |
| data_yaml_path="data.yaml", | |
| ) | |
| show_dataset_bboxes_grid(dd, 16) | |
| """ | |
| images = annotate_DetectionDataset(dataset, idxs) | |
| annotated_grid = stack_images_grid(images) | |
| sv.plot_image(annotated_grid) |