def handle(self, image, pos_points, neg_points, threshold):
image_nd = np.array(image)
clicker = Clicker()
for x, y in pos_points:
click = Click(is_positive=True, coords=(y, x))
clicker.add_click(click)
for x, y in neg_points:
click = Click(is_positive=False, coords=(y, x))
clicker.add_click(click)
predictor = get_predictor(self.net,
'NoBRS',
device=self.device,
prob_thresh=0.49)
predictor.set_input_image(image_nd)
object_prob = predictor.get_prediction(clicker)
if self.device == 'cuda':
torch.cuda.empty_cache()
object_mask = object_prob > threshold
polygon = convert_mask_to_polygon(object_mask)
return polygon
def handle(self, image, pos_points, neg_points, threshold):
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])
image_nd = input_transform(image).to(self.device)
clicker = Clicker()
for x, y in pos_points:
click = Click(is_positive=True, coords=(y, x))
clicker.add_click(click)
for x, y in neg_points:
click = Click(is_positive=False, coords=(y, x))
clicker.add_click(click)
predictor_params = {
'brs_mode': 'f-BRS-B',
'brs_opt_func_params': {'min_iou_diff': 0.001},
'lbfgs_params': {'maxfun': 20},
'predictor_params': {'max_size': 800, 'net_clicks_limit': 8},
'prob_thresh': threshold,
'zoom_in_params': {'expansion_ratio': 1.4, 'skip_clicks': 1, 'target_size': 480}}
predictor = get_predictor(self.net, device=self.device,
**predictor_params)
predictor.set_input_image(image_nd)
object_prob = predictor.get_prediction(clicker)
if self.device == 'cuda':
torch.cuda.empty_cache()
object_mask = object_prob > threshold
polygon = convert_mask_to_polygon(object_mask)
return polygon
def processing_inputs(img: Image, click_history: list, prev_polygon: list):
"""processing inputs for model
Args:
img (Image): pillow Image format
click_history (list): a list of clicks dict as [{x:x, y:y, positive:p}, ...]
prev_polygon (list): polygon shaped as (n, m, 2): [[(x, y), (x, y)...], [...], ...], where n is enclosed polygons, m is the length of coordinates
"""
img_np = np.asarray(img, dtype=np.uint8)
#gen click history
clicks = Clicker()
for c in click_history:
x = c['x']
y = c['y']
positive = c['positive']
click = Click(is_positive=positive, coords=(y, x))
clicks.add_click(click)
#gen prev mask
polygon = Polygons(prev_polygon)
mask = polygon.mask(width=img.width, height=img.height).array
prev_mask = torch.tensor(mask, dtype=torch.float32,
device=device).unsqueeze(0).unsqueeze(0)
return img_np, clicks, prev_mask
def transform(self, image_nd, clicks_lists, clicks_maps=None):
assert image_nd.shape[0] == 1 and len(clicks_lists) == 1
image_height, image_width = image_nd.shape[2:4]
self._image_crops = None
if image_height < self.crop_height or image_width < self.crop_width:
return image_nd, clicks_lists, clicks_maps
# if self._image_crops is None:
self.x_offsets = get_offsets(image_width, self.crop_width,
self.min_overlap)
self.y_offsets = get_offsets(image_height, self.crop_height,
self.min_overlap)
self._counts = np.zeros((image_height, image_width))
image_crops = []
for dy in self.y_offsets:
for dx in self.x_offsets:
self._counts[dy:dy + self.crop_height,
dx:dx + self.crop_width] += 1
image_crop = image_nd[:, :, dy:dy + self.crop_height,
dx:dx + self.crop_width]
image_crops.append(image_crop)
self._image_crops = torch.cat(*image_crops, dim=0)
self._counts = torch.tensor(self._counts,
device=image_nd.device,
dtype=torch.float32)
clicks_list = clicks_lists[0]
clicks_lists = []
for dy in self.y_offsets:
for dx in self.x_offsets:
crop_clicks = [
Click(is_positive=x.is_positive,
coords=(x.coords[0] - dy, x.coords[1] - dx))
for x in clicks_list
]
clicks_lists.append(crop_clicks)
if clicks_maps is not None:
pos_map, neg_map = clicks_maps
pos_maps_crops = []
neg_maps_crops = []
for dy in self.y_offsets:
for dx in self.x_offsets:
pos_maps_crops.append(pos_map[:, dy:dy + self.crop_height,
dx:dx + self.crop_width])
neg_maps_crops.append(neg_map[:, dy:dy + self.crop_height,
dx:dx + self.crop_width])
clicks_maps = np.concatenate(
pos_maps_crops, axis=0), np.concatenate(neg_maps_crops, axis=0)
return self._image_crops, clicks_lists, clicks_maps
def transform(self, image_nd, clicks_lists):
assert len(image_nd.shape) == 4
image_nd = torch.cat([image_nd, torch.flip(image_nd, dims=[3])], dim=0)
image_width = image_nd.shape[3]
clicks_lists_flipped = []
for clicks_list in clicks_lists:
clicks_list_flipped = [Click(is_positive=click.is_positive,
coords=(click.coords[0], image_width - click.coords[1] - 1))
for click in clicks_list]
clicks_lists_flipped.append(clicks_list_flipped)
clicks_lists = clicks_lists + clicks_lists_flipped
return image_nd, clicks_lists
def _transform_clicks(self, clicks_list):
if self._object_roi is None:
return clicks_list
rmin, rmax, cmin, cmax = self._object_roi
crop_height, crop_width = self._roi_image.shape[2:]
transformed_clicks = []
for click in clicks_list:
new_r = crop_height * (click.coords[0] - rmin) / (rmax - rmin + 1)
new_c = crop_width * (click.coords[1] - cmin) / (cmax - cmin + 1)
transformed_clicks.append(
Click(is_positive=click.is_positive, coords=(new_r, new_c)))
return transformed_clicks
def transform(self, image_nd, clicks_lists, clicks_maps=None):
assert len(image_nd.shape) == 4
image_nd = torch.cat([image_nd, torch.flip(image_nd, dims=[3])], dim=0)
image_width = image_nd.shape[3]
clicks_lists_flipped = []
for clicks_list in clicks_lists:
clicks_list_flipped = [
Click(is_positive=click.is_positive,
coords=(click.coords[0],
image_width - click.coords[1] - 1))
for click in clicks_list
]
clicks_lists_flipped.append(clicks_list_flipped)
clicks_lists = clicks_lists + clicks_lists_flipped
if clicks_maps is not None:
pos_maps, neg_maps = clicks_maps
clicks_maps = (np.concatenate((pos_maps, pos_maps[:, :, ::-1]),
axis=0),
np.concatenate((neg_maps, neg_maps[:, :, ::-1]),
axis=0))
return image_nd, clicks_lists, clicks_maps