Spaces:
Running
Running
| import numpy as np | |
| from collections import defaultdict | |
| import matplotlib.pyplot as plt | |
| import matplotlib.patches as mpatches | |
| from matplotlib import cm | |
| import torch | |
| import cv2 | |
| import random | |
| def draw_panoptic_segmentation(model,segmentation, segments_info): | |
| # get the used color map | |
| viridis = cm.get_cmap('viridis', torch.max(segmentation)) | |
| fig, ax = plt.subplots() | |
| ax.imshow(segmentation.cpu().numpy()) | |
| instances_counter = defaultdict(int) | |
| handles = [] | |
| # for each segment, draw its legend | |
| for segment in segments_info: | |
| segment_id = segment['id'] | |
| segment_label_id = segment['label_id'] | |
| segment_label = model.config.id2label[segment_label_id] | |
| label = f"{segment_label}-{instances_counter[segment_label_id]}" | |
| instances_counter[segment_label_id] += 1 | |
| color = viridis(segment_id) | |
| handles.append(mpatches.Patch(color=color, label=label)) | |
| # ax.legend(handles=handles) | |
| fig.savefig('final_mask.png') | |
| return 'final_mask.png' | |
| def draw_bboxes(rgb_frame,boxes,labels,line_thickness=3): | |
| rgb_frame = cv2.imread(rgb_frame) | |
| # rgb_frame = cv2.cvtColor(rgb_frame,cv2.COLOR_BGR2RGB) | |
| tl = line_thickness or round(0.002 * (rgb_frame.shape[0] + rgb_frame.shape[1]) / 2) + 1 # line/font thickness | |
| rgb_frame_copy = rgb_frame.copy() | |
| color_dict = {} | |
| # color = color or [random.randint(0, 255) for _ in range(3)] | |
| for item in np.unique(np.asarray(labels)): | |
| color_dict[item] = [random.randint(28, 255) for _ in range(3)] | |
| for box,label in zip(boxes,labels): | |
| if box.type() == 'torch.IntTensor': | |
| box = box.numpy() | |
| # extract coordinates | |
| x1,y1,x2,y2 = box | |
| c1,c2 = (x1,y1),(x2,y2) | |
| # Draw rectangle | |
| cv2.rectangle(rgb_frame_copy, c1,c2, color_dict[label], thickness=tl, lineType=cv2.LINE_AA) | |
| tf = max(tl - 1, 1) # font thickness | |
| # label = label2id[int(label.numpy())] | |
| t_size = cv2.getTextSize(str(label), 0, fontScale=tl / 3, thickness=tf)[0] | |
| c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 | |
| cv2.putText(rgb_frame_copy, str(label), (c1[0], c1[1] - 2), 0, tl / 3, color_dict[label], thickness=tf, lineType=cv2.LINE_AA) | |
| return rgb_frame_copy | |