fffiloni/Video-Matting-Anything
53
1import os2import cv23import torch4import logging5import numpy as np6from utils.config import CONFIG7import torch.distributed as dist8import torch.nn.functional as F9from skimage.measure import label10import pdb11 12def make_dir(target_dir):13 """14 Create dir if not exists15 """16 if not os.path.exists(target_dir):17 os.makedirs(target_dir)18 19 20def print_network(model, name):21 """22 Print out the network information23 """24 logger = logging.getLogger("Logger")25 num_params = 026 for p in model.parameters():27 num_params += p.numel()28 29 logger.info(model)30 logger.info(name)31 logger.info("Number of parameters: {}".format(num_params))32 33 34def update_lr(lr, optimizer):35 """36 update learning rates37 """38 for param_group in optimizer.param_groups:39 param_group['lr'] = lr40 41 42def warmup_lr(init_lr, step, iter_num):43 """44 Warm up learning rate45 """46 return step/iter_num*init_lr47 48 49def add_prefix_state_dict(state_dict, prefix="module"):50 """51 add prefix from the key of pretrained state dict for Data-Parallel52 """53 new_state_dict = {}54 first_state_name = list(state_dict.keys())[0]55 if not first_state_name.startswith(prefix):56 for key, value in state_dict.items():57 new_state_dict[prefix+"."+key] = state_dict[key].float()58 else:59 for key, value in state_dict.items():60 new_state_dict[key] = state_dict[key].float()61 return new_state_dict62 63 64def remove_prefix_state_dict(state_dict, prefix="module"):65 """66 remove prefix from the key of pretrained state dict for Data-Parallel67 """68 new_state_dict = {}69 first_state_name = list(state_dict.keys())[0]70 if not first_state_name.startswith(prefix):71 for key, value in state_dict.items():72 new_state_dict[key] = state_dict[key].float()73 else:74 for key, value in state_dict.items():75 new_state_dict[key[len(prefix)+1:]] = state_dict[key].float()76 return new_state_dict77 78 79def load_imagenet_pretrain(model, checkpoint_file):80 """81 Load imagenet pretrained resnet82 Add zeros channel to the first convolution layer83 Since we have the spectral normalization, we need to do a little more84 """85 checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda(CONFIG.gpu))86 state_dict = remove_prefix_state_dict(checkpoint['state_dict'])87 for key, value in state_dict.items():88 state_dict[key] = state_dict[key].float()89 90 logger = logging.getLogger("Logger")91 logger.debug("Imagenet pretrained keys:")92 logger.debug(state_dict.keys())93 logger.debug("Generator keys:")94 logger.debug(model.module.encoder.state_dict().keys())95 logger.debug("Intersection keys:")96 logger.debug(set(model.module.encoder.state_dict().keys())&set(state_dict.keys()))97 98 weight_u = state_dict["conv1.module.weight_u"]99 weight_v = state_dict["conv1.module.weight_v"]100 weight_bar = state_dict["conv1.module.weight_bar"]101 102 logger.debug("weight_v: {}".format(weight_v))103 logger.debug("weight_bar: {}".format(weight_bar.view(32, -1)))104 logger.debug("sigma: {}".format(weight_u.dot(weight_bar.view(32, -1).mv(weight_v))))105 106 new_weight_v = torch.zeros((3+CONFIG.model.mask_channel), 3, 3).cuda()107 new_weight_bar = torch.zeros(32, (3+CONFIG.model.mask_channel), 3, 3).cuda()108 109 new_weight_v[:3, :, :].copy_(weight_v.view(3, 3, 3))110 new_weight_bar[:, :3, :, :].copy_(weight_bar)111 112 logger.debug("new weight_v: {}".format(new_weight_v.view(-1)))113 logger.debug("new weight_bar: {}".format(new_weight_bar.view(32, -1)))114 logger.debug("new sigma: {}".format(weight_u.dot(new_weight_bar.view(32, -1).mv(new_weight_v.view(-1)))))115 116 state_dict["conv1.module.weight_v"] = new_weight_v.view(-1)117 state_dict["conv1.module.weight_bar"] = new_weight_bar118 119 model.module.encoder.load_state_dict(state_dict, strict=False)120 121def load_imagenet_pretrain_nomask(model, checkpoint_file):122 """123 Load imagenet pretrained resnet124 Add zeros channel to the first convolution layer125 Since we have the spectral normalization, we need to do a little more126 """127 checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda(CONFIG.gpu))128 state_dict = remove_prefix_state_dict(checkpoint['state_dict'])129 for key, value in state_dict.items():130 state_dict[key] = state_dict[key].float()131 132 logger = logging.getLogger("Logger")133 logger.debug("Imagenet pretrained keys:")134 logger.debug(state_dict.keys())135 logger.debug("Generator keys:")136 logger.debug(model.module.encoder.state_dict().keys())137 logger.debug("Intersection keys:")138 logger.debug(set(model.module.encoder.state_dict().keys())&set(state_dict.keys()))139 140 #weight_u = state_dict["conv1.module.weight_u"]141 #weight_v = state_dict["conv1.module.weight_v"]142 #weight_bar = state_dict["conv1.module.weight_bar"]143 144 #logger.debug("weight_v: {}".format(weight_v))145 #logger.debug("weight_bar: {}".format(weight_bar.view(32, -1)))146 #logger.debug("sigma: {}".format(weight_u.dot(weight_bar.view(32, -1).mv(weight_v))))147 148 #new_weight_v = torch.zeros((3+CONFIG.model.mask_channel), 3, 3).cuda()149 #new_weight_bar = torch.zeros(32, (3+CONFIG.model.mask_channel), 3, 3).cuda()150 151 #new_weight_v[:3, :, :].copy_(weight_v.view(3, 3, 3))152 #new_weight_bar[:, :3, :, :].copy_(weight_bar)153 154 #logger.debug("new weight_v: {}".format(new_weight_v.view(-1)))155 #logger.debug("new weight_bar: {}".format(new_weight_bar.view(32, -1)))156 #logger.debug("new sigma: {}".format(weight_u.dot(new_weight_bar.view(32, -1).mv(new_weight_v.view(-1)))))157 158 #state_dict["conv1.module.weight_v"] = new_weight_v.view(-1)159 #state_dict["conv1.module.weight_bar"] = new_weight_bar160 161 model.module.encoder.load_state_dict(state_dict, strict=False)162 163def load_VGG_pretrain(model, checkpoint_file):164 """165 Load imagenet pretrained resnet166 Add zeros channel to the first convolution layer167 Since we have the spectral normalization, we need to do a little more168 """169 checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda())170 backbone_state_dict = remove_prefix_state_dict(checkpoint['state_dict'])171 172 model.module.encoder.load_state_dict(backbone_state_dict, strict=False)173 174 175def get_unknown_tensor(trimap):176 """177 get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor178 """179 if trimap.shape[1] == 3:180 weight = trimap[:, 1:2, :, :].float()181 else:182 weight = trimap.eq(1).float()183 return weight184 185def get_gaborfilter(angles):186 """187 generate gabor filter as the conv kernel188 :param angles: number of different angles189 """190 gabor_filter = []191 for angle in range(angles):192 gabor_filter.append(cv2.getGaborKernel(ksize=(5,5), sigma=0.5, theta=angle*np.pi/8, lambd=5, gamma=0.5))193 gabor_filter = np.array(gabor_filter)194 gabor_filter = np.expand_dims(gabor_filter, axis=1)195 return gabor_filter.astype(np.float32)196 197 198def get_gradfilter():199 """200 generate gradient filter as the conv kernel201 """202 grad_filter = []203 grad_filter.append([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])204 grad_filter.append([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])205 grad_filter = np.array(grad_filter)206 grad_filter = np.expand_dims(grad_filter, axis=1)207 return grad_filter.astype(np.float32)208 209 210def reduce_tensor_dict(tensor_dict, mode='mean'):211 """212 average tensor dict over different GPUs213 """214 for key, tensor in tensor_dict.items():215 if tensor is not None:216 tensor_dict[key] = reduce_tensor(tensor, mode)217 return tensor_dict218 219 220def reduce_tensor(tensor, mode='mean'):221 """222 average tensor over different GPUs223 """224 rt = tensor.clone()225 dist.all_reduce(rt, op=dist.ReduceOp.SUM)226 if mode == 'mean':227 rt /= CONFIG.world_size228 elif mode == 'sum':229 pass230 else:231 raise NotImplementedError("reduce mode can only be 'mean' or 'sum'")232 return rt233 234### preprocess the image and mask for inference (np array), crop based on ROI235def preprocess(image, mask, thres):236 mask_ = (mask >= thres).astype(np.float32)237 arr = np.nonzero(mask_)238 h, w = mask.shape239 bbox = [max(0, int(min(arr[0]) - 0.1*h)),240 min(h, int(max(arr[0]) + 0.1*h)),241 max(0, int(min(arr[1]) - 0.1*w)),242 min(w, int(max(arr[1]) + 0.1*w))]243 image = image[bbox[0]:bbox[1], bbox[2]:bbox[3], :]244 mask = mask[bbox[0]:bbox[1], bbox[2]:bbox[3]]245 return image, mask, bbox246 247### postprocess the alpha prediction to keep the largest connected component (np array) and uncrop, alpha in [0, 1]248### based on https://github.com/senguptaumd/Background-Matting/blob/master/test_background-matting_image.py249def postprocess(alpha, orih=None, oriw=None, bbox=None):250 labels=label((alpha>0.05).astype(int))251 try:252 assert( labels.max() != 0 )253 except:254 return None255 largestCC = labels == np.argmax(np.bincount(labels.flat)[1:])+1256 alpha = alpha * largestCC257 if bbox is None:258 return alpha259 else:260 ori_alpha = np.zeros(shape=[orih, oriw], dtype=np.float32)261 ori_alpha[bbox[0]:bbox[1], bbox[2]:bbox[3]] = alpha262 return ori_alpha263 264 265Kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,30)]266def get_unknown_tensor_from_pred(pred, rand_width=30, train_mode=True):267 ### pred: N, 1 ,H, W 268 N, C, H, W = pred.shape269 270 pred = pred.data.cpu().numpy()271 uncertain_area = np.ones_like(pred, dtype=np.uint8)272 uncertain_area[pred<1.0/255.0] = 0273 uncertain_area[pred>1-1.0/255.0] = 0274 for n in range(N):275 uncertain_area_ = uncertain_area[n,0,:,:] # H, W276 if train_mode:277 width = np.random.randint(1, rand_width)278 else:279 width = rand_width // 2280 uncertain_area_ = cv2.dilate(uncertain_area_, Kernels[width])281 uncertain_area[n,0,:,:] = uncertain_area_282 weight = np.zeros_like(uncertain_area)283 weight[uncertain_area == 1] = 1284 weight = torch.from_numpy(weight).cuda()285 286 return weight287 288def get_unknown_tensor_from_pred_oneside(pred, rand_width=30, train_mode=True):289 ### pred: N, 1 ,H, W 290 N, C, H, W = pred.shape291 pred = pred.data.cpu().numpy()292 uncertain_area = np.ones_like(pred, dtype=np.uint8)293 uncertain_area[pred<1.0/255.0] = 0294 #uncertain_area[pred>1-1.0/255.0] = 0295 for n in range(N):296 uncertain_area_ = uncertain_area[n,0,:,:] # H, W297 if train_mode:298 width = np.random.randint(1, rand_width)299 else:300 width = rand_width // 2301 uncertain_area_ = cv2.dilate(uncertain_area_, Kernels[width])302 uncertain_area[n,0,:,:] = uncertain_area_303 uncertain_area[pred>1-1.0/255.0] = 0304 #weight = np.zeros_like(uncertain_area)305 #weight[uncertain_area == 1] = 1306 weight = torch.from_numpy(uncertain_area).cuda()307 return weight308 309Kernels_mask = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,30)]310def get_unknown_tensor_from_mask(mask, rand_width=30, train_mode=True):311 """312 get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor313 """314 N, C, H, W = mask.shape315 mask_c = mask.data.cpu().numpy().astype(np.uint8)316 317 weight = np.ones_like(mask_c, dtype=np.uint8)318 319 for n in range(N):320 if train_mode:321 width = np.random.randint(rand_width // 2, rand_width)322 else:323 width = rand_width // 2324 fg_mask = cv2.erode(mask_c[n,0], Kernels_mask[width])325 bg_mask = cv2.erode(1 - mask_c[n,0], Kernels_mask[width])326 weight[n,0][fg_mask==1] = 0327 weight[n,0][bg_mask==1] = 0328 weight = torch.from_numpy(weight).cuda()329 return weight330 331def get_unknown_tensor_from_mask_oneside(mask, rand_width=30, train_mode=True):332 """333 get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor334 """335 N, C, H, W = mask.shape336 mask_c = mask.data.cpu().numpy().astype(np.uint8)337 338 weight = np.ones_like(mask_c, dtype=np.uint8)339 340 for n in range(N):341 if train_mode:342 width = np.random.randint(rand_width // 2, rand_width)343 else:344 width = rand_width // 2345 #fg_mask = cv2.erode(mask_c[n,0], Kernels_mask[width])346 fg_mask = mask_c[n,0]347 bg_mask = cv2.erode(1 - mask_c[n,0], Kernels_mask[width])348 weight[n,0][fg_mask==1] = 0349 weight[n,0][bg_mask==1] = 0350 weight = torch.from_numpy(weight).cuda()351 return weight352 353def get_unknown_box_from_mask(mask):354 """355 get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor356 """357 N, C, H, W = mask.shape358 mask_c = mask.data.cpu().numpy().astype(np.uint8)359 360 weight = np.ones_like(mask_c, dtype=np.uint8)361 fg_set = np.where(mask_c[0][0] != 0)362 x_min = np.min(fg_set[1])363 x_max = np.max(fg_set[1])364 y_min = np.min(fg_set[0])365 y_max = np.max(fg_set[0])366 367 weight[0, 0, y_min:y_max, x_min:x_max] = 0368 weight = torch.from_numpy(weight).cuda()369 return weight