Spaces:
Sleeping
Sleeping
File size: 2,135 Bytes
0e78cbf d766b17 0e78cbf d766b17 |
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 |
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,color=None,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()
if color is None :
color = color or [random.randint(0, 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, 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, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
return rgb_frame_copy |