def get_visualization(self, alpha_blend, click_radius):
if self.image is None:
return None
results_mask_for_vis = self.result_mask
if self.probs_history:
results_mask_for_vis[self.current_object_prob > self.prob_thresh] = self.object_count + 1
vis = draw_with_blend_and_clicks(self.image, mask=results_mask_for_vis, alpha=alpha_blend,
clicks_list=self.clicker.clicks_list, radius=click_radius)
if self.probs_history:
total_mask = self.probs_history[-1][0] > self.prob_thresh
results_mask_for_vis[np.logical_not(total_mask)] = 0
vis = draw_with_blend_and_clicks(vis, mask=results_mask_for_vis, alpha=alpha_blend)
return vis
def prepare_result(img: np.ndarray,
pred_probs: np.ndarray,
clicks: Clicker,
gt_mask_file,
tolerance: int = 1,
view_img: bool = False,
filename: str = None):
"""prepare result
Args:
img (np.ndarray): img in numpy.ndarray
pred_mask (np.ndarray): predicted mask from model
clicks(Clicker): Cliker object for click history
gt_mask_file (FileStorage): ground truth mask file
tolerance (int, optional): Precision to convert from mask to polygon in pixel. Defaults to 1.
view_img (bool, optional): Return result image url. Defaults to False.
"""
# gen mask
assert len(
pred_probs
) == 1, f'Only one output is expected, but got {len(pred_probs)}'
pred_probs = pred_probs[0]
pred_mask = pred_probs > MODEL_THRESH
# convert mask to polygon
regions = Mask(pred_mask).polygons().points
polygons = []
for polygon in regions:
polygon2 = measure.approximate_polygon(polygon, tolerance)
polygons.append(polygon2.tolist())
results = {'polygons': polygons}
# calculate iou
if gt_mask_file:
gt_mask = Image.open(gt_mask_file)
mask_np = np.asarray(gt_mask, dtype=np.int32)
if len(mask_np.shape) > 2:
assert len(mask_np.shape) == 3
mask_np = np.max(mask_np, axis=2)
mask_np = np.where(mask_np > 0, 1, 0)
iou = utils.get_iou(mask_np, pred_mask)
results['iou'] = iou
print(iou)
# save img with minor delay
if view_img:
ext = filename.split('.')[-1]
draw = vis.draw_with_blend_and_clicks(img,
mask=pred_mask,
clicks_list=clicks.clicks_list)
filename = filename.split('.')[0] + f'[{len(clicks.clicks_list)}].jpg'
result_path = TEMP_PATH + filename
Image.fromarray(draw).save(result_path)
# return send_file(result_path)
results['result'] = filename
return results
def callback(image, gt_mask, pred_probs, sample_id, click_indx,
clicks_list):
sample_path = save_path / f'{sample_id}_{click_indx}.jpg'
prob_map = draw_probmap(pred_probs)
image_with_mask = draw_with_blend_and_clicks(image,
pred_probs > prob_thresh,
clicks_list=clicks_list)
cv2.imwrite(
str(sample_path),
np.concatenate((image_with_mask, prob_map), axis=1)[:, :, ::-1])
def _update_image(self, reset_canvas=False):
alpha = self.state['alpha_blend'].get()
click_radius = self.state['click_radius'].get()
mask = self._get_current_mask()
image = vis.draw_with_blend_and_clicks(self.state['_image'], mask=mask, alpha=alpha,
clicks_list=self.state['_clicks_list'], radius=click_radius)
if self.image_on_canvas is None:
self.image_on_canvas = CanvasImage(self.canvas_frame, self.canvas)
self.image_on_canvas.register_click_callback(self._click_callback)
self.image_on_canvas.reload_image(Image.fromarray(image), reset_canvas)
brs_mode = 'f-BRS-B'
predictor = get_predictor(model, brs_mode, device, prob_thresh=MODEL_THRESH)
## single sample test
sample_id = random.sample(range(len(dataset)), 1)[0]
TARGET_IOU = 0.95
sample = dataset.get_sample(sample_id)
gt_mask = sample.gt_mask
clicks_list, ious_arr, pred = evaluate_sample(sample.image,
gt_mask,
predictor,
pred_thr=MODEL_THRESH,
max_iou_thr=TARGET_IOU,
max_clicks=EVAL_MAX_CLICKS)
pred_mask = pred > MODEL_THRESH
draw = vis.draw_with_blend_and_clicks(sample.image,
mask=pred_mask,
clicks_list=clicks_list)
draw = np.concatenate(
(draw, 255 * pred_mask[:, :, np.newaxis].repeat(3, axis=2), 255 *
(gt_mask > 0)[:, :, np.newaxis].repeat(3, axis=2)),
axis=1)
print(ious_arr)
plt.figure(figsize=(20, 30))
plt.imshow(draw)
plt.show()