File size: 1,777 Bytes
0e83290 |
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 |
from .fbrs.controller import InteractiveController
from .fbrs.inference import utils
class FBRSController:
def __init__(self, checkpoint_path, device='cuda:0', max_size=800):
model = utils.load_is_model(checkpoint_path, device, cpu_dist_maps=True, norm_radius=260)
# Predictor params
zoomin_params = {
'skip_clicks': 1,
'target_size': 480,
'expansion_ratio': 1.4,
}
predictor_params = {
'brs_mode': 'f-BRS-B',
'prob_thresh': 0.5,
'zoom_in_params': zoomin_params,
'predictor_params': {
'net_clicks_limit': 8,
'max_size': 800,
},
'brs_opt_func_params': {'min_iou_diff': 1e-3},
'lbfgs_params': {'maxfun': 20}
}
self.controller = InteractiveController(model, device, predictor_params)
self.anchored = False
self.device = device
def unanchor(self):
self.anchored = False
def interact(self, image, x, y, is_positive):
image = image.to(self.device, non_blocking=True)
if not self.anchored:
self.controller.set_image(image)
self.controller.reset_predictor()
self.anchored = True
self.controller.add_click(x, y, is_positive)
# return self.controller.result_mask
# return self.controller.probs_history[-1][1]
return (self.controller.probs_history[-1][1]>0.5).float()
def undo(self):
self.controller.undo_click()
if len(self.controller.probs_history) == 0:
return None
else:
return (self.controller.probs_history[-1][1]>0.5).float() |