def detect_frame(capture, model, label):
"""Use the learner to detect objects.
"""
_, frame = capture.read() # Capture frame-by-frame
preds = model([utils.cv2torch(frame)])
anno_bboxes = _get_det_bboxes(preds, labels=label)
im_pil = utils.cv2pil(frame)
plot_boxes(im_pil,
anno_bboxes,
plot_settings=PlotSettings(rect_color=(0, 255, 0)))
return utils.pil2matplotlib(im_pil)
def test_plot_mask(od_mask_rects):
""" Test that `plot_mask` works. """
plot_setting = PlotSettings()
_, mask, rects, im = od_mask_rects
# plot mask
im = plot_mask(im, mask, plot_settings=plot_setting).convert('RGB')
im = np.transpose(np.array(im), (2, 0, 1))
# validate each channel matches the mask
for ch in im:
ch_uniques = np.unique(ch)
foreground_uniques = np.unique(ch[np.where(mask != 0)])
assert len(foreground_uniques) == 1
assert foreground_uniques[0] == ch_uniques[1]
background_uniques = np.unique(ch[np.where(mask == 0)])
assert len(background_uniques) == 1
assert background_uniques[0] == ch_uniques[0]
def basic_plot_settings() -> PlotSettings:
return PlotSettings()