def overlay_boxes(self, image, predictions):
"""
Adds the predicted boxes on top of the image
Arguments:
image (np.ndarray): an image as returned by OpenCV
predictions (BoxList): the result of the computation by the model.
It should contain the field `labels`.
"""
labels = predictions.get_field("labels")
boxes = predictions.bbox
# colors = self.compute_colors_for_labels(labels).tolist()
colors = colormap(rgb=True).tolist()
# colors = self.compute_colors_for_labels_yolact(labels)
for box, color in zip(boxes, colors):
box = box.to(torch.int64)
top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
image = cv2.rectangle(
image, tuple(top_left), tuple(bottom_right), tuple(color), 1
)
# cv2.imshow('test',image)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
return image
def overlay_mask(self,
image,
predictions,
alpha=0.3,
border_alpha=1.0,
border_thick=2):
"""
Adds the instances contours for each predicted object.
Each label has a different color.
Arguments:
image (np.ndarray): an image as returned by OpenCV
predictions (BoxList): the result of the computation by the model.
It should contain the field `mask` and `labels`.
"""
masks = predictions.get_field("mask").numpy()
labels = predictions.get_field("labels")
colors = colormap(rgb=True)[:len(labels)]
for mask, color in zip(masks, colors):
image = image.astype(np.float32)
mask = mask[0, :, :, None]
idx = np.nonzero(mask)
image[idx[0], idx[1], :] *= 1.0 - alpha
image[idx[0], idx[1], :] += [alpha * x for x in color]
if border_alpha == 0:
continue
border_color = [x * 0.5 for x in color]
if isinstance(border_color, np.ndarray):
border_color = border_color.tolist()
contours, _ = cv2_util.findContours(mask, cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
if border_alpha < 1:
with_border = image.copy()
cv2.drawContours(with_border, contours, -1, border_color,
border_thick, cv2.LINE_AA)
image = ((1 - border_alpha) * image +
border_alpha * with_border)
else:
cv2.drawContours(image, contours, -1, border_color,
border_thick, cv2.LINE_AA)
return image
def overlay_class_names(self, image, predictions, display_score=False):
"""
Adds detected class names and scores in the positions defined by the
top-left corner of the predicted bounding box
Arguments:
image (np.ndarray): an image as returned by OpenCV
predictions (BoxList): the result of the computation by the model.
It should contain the field `scores` and `labels`.
"""
scores = predictions.get_field("scores").tolist()
# labels = predictions.get_field("labels").tolist()
labels = predictions.get_field("labels")
# colors = self.compute_colors_for_labels(labels).tolist()
colors = colormap(rgb=True).tolist()
# colors = self.compute_colors_for_labels_yolact(labels)
labels = labels.tolist()
labels = [self.CATEGORIES[i] for i in labels]
boxes = predictions.bbox
# font_face = cv2.FONT_HERSHEY_COMPLEX
font_face = cv2.FONT_HERSHEY_DUPLEX
font_scale = 0.6
font_thickness = 1
for box, score, label, color in zip(boxes, scores, labels, colors):
x, y = box[:2]
if display_score:
template = "{}: {:.2f}"
s = template.format(label, score)
else:
s = label
text_w, text_h = cv2.getTextSize(s, font_face, font_scale,
font_thickness)[0]
text_pt = (x, y - 3)
text_color = [255, 255, 255]
# cv2.rectangle(image, (x,y), (x + text_w, y - text_h - 2), color, -1)
# cv2.putText(image, s, (x, y), font_face, font_scale, (255, 255, 255), 1)
cv2.rectangle(image, (x, y), (x + text_w, y - text_h - 4), color,
-1) # mimicing yolact
cv2.putText(image, s, text_pt, font_face, font_scale, text_color,
font_thickness, cv2.LINE_AA)
return image
def overlay_mask(self, image, predictions):
"""
Adds the instances FCOS-MASK-SCORING-PLUS-R-50-FPN-MaskP3P5-new2xcontours for each predicted object.
Each label has a different color.
Arguments:
image (np.ndarray): an image as returned by OpenCV
predictions (BoxList): the result of the computation by the model.
It should contain the field `mask` and `labels`.
"""
masks = predictions.get_field("mask").numpy()
labels = predictions.get_field("labels")
#original
# colors = self.compute_colors_for_labels(labels).tolist()
#Detectron.pytorch for matplotlib colors
colors = colormap(rgb=True).tolist()
# colors = self.compute_colors_for_labels_yolact(labels)
mask_img = np.copy(image)
for mask, color in zip(masks, colors):
# color_mask = color_list[color_id % len(color_list)]
# color_id += 1
thresh = mask[0, :, :, None]
contours, hierarchy = cv2_util.findContours(
thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
)
mask_img = cv2.drawContours(mask_img, contours, -1, color, -1)
# composite = image
alpha = 0.45
composite = cv2.addWeighted(image, 1.0 - alpha, mask_img, alpha, 0)
# composite = cv2.addWeighted(image, 1.0 - alpha, mask, alpha, 0)
return composite