Spaces:
Runtime error
Runtime error
| """ | |
| Contains all the types of interaction related to the GUI | |
| Not related to automatic evaluation in the DAVIS dataset | |
| You can inherit the Interaction class to create new interaction types | |
| undo is (sometimes partially) supported | |
| """ | |
| import torch | |
| import numpy as np | |
| import cv2 | |
| import time | |
| from PIL import Image | |
| from ..model.aggregate import aggregate_sbg, aggregate_wbg | |
| from ..util.tensor_util import pad_divide_by, unpad | |
| from collections import deque | |
| from copy import deepcopy | |
| from .interactive_utils import color_map | |
| max_history = 50 | |
| class Interaction: | |
| def __init__(self, image, prev_mask, true_size, controller): | |
| self.image = image # This image is PADDED | |
| self.prev_mask = prev_mask.clone() # This is also PADDED | |
| self.controller = controller | |
| self.start_time = time.time() | |
| self.history = deque(maxlen=max_history) | |
| self.h, self.w = true_size | |
| self.out_prob = None | |
| self.out_mask = None | |
| def undo(self): | |
| pass | |
| def can_undo(self): | |
| return len(self.history) > 0 | |
| def predict(self): | |
| pass | |
| class LocalInteraction(Interaction): | |
| # This interaction compress all the interactions done in the local stage into one | |
| # Performs stitching | |
| def __init__(self, image, prev_mask, true_size, bounding_box, region_prob, pad, local_pad): | |
| super().__init__(image, prev_mask, true_size, None) | |
| lx, ux, ly, uy = bounding_box | |
| self.out_prob = unpad(self.prev_mask, pad) | |
| region_prob = unpad(region_prob, local_pad) | |
| # Trim the margin since results at the boundary are not that trustworthy | |
| if (ux-lx) > 6 and (uy-ly) > 6: | |
| lx += 3 | |
| ux -= 3 | |
| ly += 3 | |
| uy -= 3 | |
| self.out_prob[:,:,ly:uy+1, lx:ux+1] = region_prob[:,:,3:-3,3:-3] | |
| else: | |
| self.out_prob[:,:,ly:uy+1, lx:ux+1] = region_prob | |
| self.out_prob, _ = pad_divide_by(self.out_prob, 16, self.out_prob.shape[-2:]) | |
| self.out_mask = aggregate_sbg(self.out_prob, keep_bg=True) | |
| self.storage = None # Might be used outside | |
| def can_undo(self): | |
| return False | |
| def predict(self): | |
| return self.out_mask | |
| class CropperInteraction(Interaction): | |
| # Turns a global map into a local map through cropping | |
| def __init__(self, image, prev_mask, pad, bounding_box): | |
| lx, ux, ly, uy = bounding_box | |
| true_size = (uy-ly+1, ux-lx+1) | |
| super().__init__(image, prev_mask, true_size, None) | |
| self.bounding_box = bounding_box # UN-PADDED | |
| unpad_prev_mask = unpad(self.prev_mask, pad) | |
| self.out_prob = unpad_prev_mask[:, :, ly:uy+1, lx:ux+1] | |
| self.out_prob, self.pad = pad_divide_by(self.out_prob, 16, self.out_prob.shape[-2:]) | |
| self.out_mask = aggregate_sbg(self.out_prob, keep_bg=True) | |
| unpad_image = unpad(self.image, pad) | |
| self.im_crop = unpad_image[:, :, ly:uy+1, lx:ux+1] | |
| self.im_crop, _ = pad_divide_by(self.im_crop, 16, self.im_crop.shape[-2:]) | |
| def can_undo(self): | |
| return False | |
| def predict(self): | |
| return self.out_mask | |
| class FreeInteraction(Interaction): | |
| def __init__(self, image, prev_mask, true_size, num_objects, pad): | |
| """ | |
| prev_mask should be in probabilities | |
| """ | |
| super().__init__(image, prev_mask, true_size, None) | |
| self.K = num_objects | |
| self.drawn_map = unpad(self.prev_mask, pad).detach().cpu().numpy() | |
| self.curr_path = [[] for _ in range(self.K + 1)] | |
| self.all_paths = [self.curr_path] | |
| self.size = None | |
| self.surplus_history = False | |
| def set_size(self, size): | |
| self.size = size | |
| """ | |
| k - object id | |
| vis - a tuple (visualization map, pass through alpha). None if not needed. | |
| """ | |
| def push_point(self, x, y, k, vis=None): | |
| if vis is not None: | |
| vis_map, vis_alpha = vis | |
| selected = self.curr_path[k] | |
| selected.append((x, y)) | |
| if len(selected) >= 2: | |
| for i in range(self.K): | |
| self.drawn_map[i,0] = cv2.line(self.drawn_map[i,0], | |
| (int(round(selected[-2][0])), int(round(selected[-2][1]))), | |
| (int(round(selected[-1][0])), int(round(selected[-1][1]))), | |
| int((i+1)==k), thickness=self.size) | |
| # Plot visualization | |
| if vis is not None: | |
| # Visualization for drawing | |
| if k == 0: | |
| vis_map = cv2.line(vis_map, | |
| (int(round(selected[-2][0])), int(round(selected[-2][1]))), | |
| (int(round(selected[-1][0])), int(round(selected[-1][1]))), | |
| color_map[k], thickness=self.size) | |
| else: | |
| vis_map = cv2.line(vis_map, | |
| (int(round(selected[-2][0])), int(round(selected[-2][1]))), | |
| (int(round(selected[-1][0])), int(round(selected[-1][1]))), | |
| color_map[k], thickness=self.size) | |
| # Visualization on/off boolean filter | |
| vis_alpha = cv2.line(vis_alpha, | |
| (int(round(selected[-2][0])), int(round(selected[-2][1]))), | |
| (int(round(selected[-1][0])), int(round(selected[-1][1]))), | |
| 0.75, thickness=self.size) | |
| if vis is not None: | |
| return vis_map, vis_alpha | |
| def end_path(self): | |
| # Complete the drawing | |
| self.curr_path = [[] for _ in range(self.K + 1)] | |
| self.all_paths.append(self.curr_path) | |
| self.history.append(self.drawn_map.copy()) | |
| self.surplus_history = True | |
| def predict(self): | |
| self.out_prob = torch.from_numpy(self.drawn_map).float().cuda() | |
| self.out_prob, _ = pad_divide_by(self.out_prob, 16, self.out_prob.shape[-2:]) | |
| self.out_mask = aggregate_sbg(self.out_prob, keep_bg=True) | |
| return self.out_mask | |
| def undo(self): | |
| if self.surplus_history: | |
| self.history.pop() | |
| self.surplus_history = False | |
| self.drawn_map = self.history.pop() | |
| # pop the current path (which is empty) and the last path | |
| self.all_paths = self.all_paths[:-2] | |
| self.curr_path = [[] for _ in range(self.K + 1)] | |
| self.all_paths.append(self.curr_path) | |
| return self.predict() | |
| def can_undo(self): | |
| return (len(self.history) > 0) and not ( | |
| self.surplus_history and (len(self.history) < 2)) | |
| class ScribbleInteraction(Interaction): | |
| def __init__(self, image, prev_mask, true_size, controller, num_objects): | |
| """ | |
| prev_mask should be in an indexed form | |
| """ | |
| super().__init__(image, prev_mask, true_size, controller) | |
| self.K = num_objects | |
| self.drawn_map = np.empty((self.h, self.w), dtype=np.uint8) | |
| self.drawn_map.fill(255) | |
| # background + k | |
| self.curr_path = [[] for _ in range(self.K + 1)] | |
| self.all_paths = [self.curr_path] | |
| self.size = 3 | |
| self.surplus_history = False | |
| """ | |
| k - object id | |
| vis - a tuple (visualization map, pass through alpha). None if not needed. | |
| """ | |
| def push_point(self, x, y, k, vis=None): | |
| if vis is not None: | |
| vis_map, vis_alpha = vis | |
| selected = self.curr_path[k] | |
| selected.append((x, y)) | |
| if len(selected) >= 2: | |
| self.drawn_map = cv2.line(self.drawn_map, | |
| (int(round(selected[-2][0])), int(round(selected[-2][1]))), | |
| (int(round(selected[-1][0])), int(round(selected[-1][1]))), | |
| k, thickness=self.size) | |
| # Plot visualization | |
| if vis is not None: | |
| # Visualization for drawing | |
| if k == 0: | |
| vis_map = cv2.line(vis_map, | |
| (int(round(selected[-2][0])), int(round(selected[-2][1]))), | |
| (int(round(selected[-1][0])), int(round(selected[-1][1]))), | |
| color_map[k], thickness=self.size) | |
| else: | |
| vis_map = cv2.line(vis_map, | |
| (int(round(selected[-2][0])), int(round(selected[-2][1]))), | |
| (int(round(selected[-1][0])), int(round(selected[-1][1]))), | |
| color_map[k], thickness=self.size) | |
| # Visualization on/off boolean filter | |
| vis_alpha = cv2.line(vis_alpha, | |
| (int(round(selected[-2][0])), int(round(selected[-2][1]))), | |
| (int(round(selected[-1][0])), int(round(selected[-1][1]))), | |
| 0.75, thickness=self.size) | |
| # Optional vis return | |
| if vis is not None: | |
| return vis_map, vis_alpha | |
| def end_path(self): | |
| # Complete the drawing | |
| self.curr_path = [[] for _ in range(self.K + 1)] | |
| self.all_paths.append(self.curr_path) | |
| self.history.append(self.drawn_map.copy()) | |
| self.surplus_history = True | |
| def predict(self): | |
| self.out_prob = self.controller.interact(self.image, self.prev_mask, self.drawn_map) | |
| self.out_mask = aggregate_wbg(self.out_prob, keep_bg=True, hard=True) | |
| return self.out_mask | |
| def undo(self): | |
| if self.surplus_history: | |
| self.history.pop() | |
| self.surplus_history = False | |
| self.drawn_map = self.history.pop() | |
| # pop the current path (which is empty) and the last path | |
| self.all_paths = self.all_paths[:-2] | |
| self.curr_path = [[] for _ in range(self.K + 1)] | |
| self.all_paths.append(self.curr_path) | |
| return self.predict() | |
| def can_undo(self): | |
| return (len(self.history) > 0) and not ( | |
| self.surplus_history and (len(self.history) < 2)) | |
| class ClickInteraction(Interaction): | |
| def __init__(self, image, prev_mask, true_size, controller, tar_obj, pad): | |
| """ | |
| prev_mask in a prob. form | |
| """ | |
| super().__init__(image, prev_mask, true_size, controller) | |
| self.tar_obj = tar_obj | |
| self.pad = pad | |
| # negative/positive for each object | |
| self.pos_clicks = [] | |
| self.neg_clicks = [] | |
| self.out_prob = self.prev_mask.clone() | |
| self.surplus_history = False | |
| """ | |
| neg - Negative interaction or not | |
| vis - a tuple (visualization map, pass through alpha). None if not needed. | |
| """ | |
| def push_point(self, x, y, neg, vis=None): | |
| # Clicks | |
| if neg: | |
| self.neg_clicks.append((x, y)) | |
| else: | |
| self.pos_clicks.append((x, y)) | |
| # Do the prediction, note that the image is padded | |
| self.obj_mask = self.controller.interact(self.image, x+self.pad[0], y+self.pad[2], not neg) | |
| self.history.append(deepcopy((self.pos_clicks, self.neg_clicks))) | |
| self.surplus_history = True | |
| # Plot visualization | |
| if vis is not None: | |
| vis_map, vis_alpha = vis | |
| # Visualization for clicks | |
| if neg: | |
| vis_map = cv2.circle(vis_map, | |
| (int(round(x)), int(round(y))), | |
| 2, color_map[0], thickness=-1) | |
| else: | |
| vis_map = cv2.circle(vis_map, | |
| (int(round(x)), int(round(y))), | |
| 2, color_map[self.tar_obj], thickness=-1) | |
| vis_alpha = cv2.circle(vis_alpha, | |
| (int(round(x)), int(round(y))), | |
| 2, 1, thickness=-1) | |
| # Optional vis return | |
| return vis_map, vis_alpha | |
| def predict(self): | |
| if self.obj_mask is None: | |
| self.out_prob = self.prev_mask.clone() | |
| else: | |
| self.out_prob[self.tar_obj-1] = self.obj_mask | |
| self.out_mask = aggregate_sbg(self.out_prob, keep_bg=True, hard=True) | |
| return self.out_mask | |
| def undo(self): | |
| if self.surplus_history: | |
| self.history.pop() | |
| self.surplus_history = False | |
| self.pos_clicks, self.neg_clicks = self.history.pop() | |
| self.obj_mask = self.controller.undo() | |
| return self.predict() | |
| def can_undo(self): | |
| return (len(self.history) > 0) and not ( | |
| self.surplus_history and (len(self.history) < 2)) |