Spaces:
Running
Running
File size: 3,673 Bytes
dd90166 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | 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) |