def draw_rect(self, image, pred, segms=True):
"""segms=False 不画mask"""
labels = pred["labels"]
bboxs = pred["boxes"]
scores = pred["scores"]
if segms and self.useMask:
masks = pred["masks"]
mask_color_id = 0
color_list = colormap.colormap()
for idx, (label, bbox, score) in enumerate(zip(labels, bboxs, scores)):
label = label.cpu().numpy()
bbox = bbox.cpu().numpy() # .astype(np.int16)
score = score.cpu().numpy()
class_str = "%s:%.3f" % (self.classes[int(label) - 1], score
) # 跳过背景
pos = list(map(int, bbox))
image = opencv.vis_rect(image,
pos,
class_str,
0.5,
int(label),
useMask=False if self.useMask else True)
if segms and self.useMask:
mask = masks[idx].cpu().numpy()
mask = np.clip(np.squeeze(mask, 0) * 255., 0,
255).astype(np.uint8)
mask_color_id += 1
color_mask = color_list[mask_color_id % len(color_list), 0:3]
image = opencv.vis_mask(image, mask, color_mask, 0.3, True)
return image
def draw_rect(self, image, pred):
labels = pred["labels"]
bboxs = pred["boxes"]
scores = pred["scores"]
for label, bbox, score in zip(labels, bboxs, scores):
label = label.cpu().numpy()
bbox = bbox.cpu().numpy() #.astype(np.int16)
score = score.cpu().numpy()
class_str = "%s:%.3f" % (self.classes[int(label)], score) # 跳过背景
pos = list(map(int, bbox))
image = opencv.vis_rect(image, pos, class_str, 0.5, int(label))
return image