def val(self): """ Validation function during the train phase. """ self.cls_net.eval() start_time = time.time() with torch.no_grad(): for j, data_dict in enumerate(self.val_loader): # Forward pass. data_dict = RunnerHelper.to_device(self, data_dict) out = self.cls_net(data_dict) loss_dict = self.loss(out) out_dict, label_dict, _ = RunnerHelper.gather(self, out) self.running_score.update(out_dict, label_dict) self.val_losses.update({key: loss.item() for key, loss in loss_dict.items()}, data_dict['img'].size(0)) # Update the vars of the val phase. self.batch_time.update(time.time() - start_time) start_time = time.time() RunnerHelper.save_net(self, self.cls_net) # Print the log info & reset the states. Log.info('Test Time {batch_time.sum:.3f}s'.format(batch_time=self.batch_time)) Log.info('TestLoss = {}'.format(self.val_losses.info())) Log.info('Top1 ACC = {}'.format(RunnerHelper.dist_avg(self, self.running_score.get_top1_acc()))) Log.info('Top3 ACC = {}'.format(RunnerHelper.dist_avg(self, self.running_score.get_top3_acc()))) Log.info('Top5 ACC = {}'.format(RunnerHelper.dist_avg(self, self.running_score.get_top5_acc()))) self.batch_time.reset() self.batch_time.reset() self.val_losses.reset() self.running_score.reset() self.cls_net.train()
def __list_dirs(self, root_dir, dataset): img_list = list() label_list = list() image_dir = os.path.join(root_dir, dataset, 'image') label_dir = os.path.join(root_dir, dataset, 'label') for file_name in os.listdir(label_dir): image_name = '.'.join(file_name.split('.')[:-1]) label_path = os.path.join(label_dir, file_name) img_path = ImageHelper.imgpath(image_dir, image_name) if not os.path.exists(label_path) or img_path is None: Log.warn('Label Path: {} not exists.'.format(label_path)) continue img_list.append(img_path) label_list.append(label_path) if dataset == 'train' and self.configer.get('data', 'include_val'): image_dir = os.path.join(root_dir, 'val/image') label_dir = os.path.join(root_dir, 'val/label') for file_name in os.listdir(label_dir): image_name = '.'.join(file_name.split('.')[:-1]) label_path = os.path.join(label_dir, file_name) img_path = ImageHelper.imgpath(image_dir, image_name) if not os.path.exists(label_path) or img_path is None: Log.warn('Label Path: {} not exists.'.format(label_path)) continue img_list.append(img_path) label_list.append(label_path) return img_list, label_list
def __call__(self): # the 30th layer of features is relu of conv5_3 model = vgg16(pretrained=False) if self.configer.get('network', 'pretrained') is not None: Log.info('Loading pretrained model: {}'.format( self.configer.get('network', 'pretrained'))) model.load_state_dict( torch.load(self.configer.get('network', 'pretrained'))) features = list(model.features)[:30] classifier = model.classifier classifier = list(classifier) del classifier[6] if not self.configer.get('network', 'use_drop'): del classifier[5] del classifier[2] classifier = nn.Sequential(*classifier) # freeze top4 conv for layer in features[:10]: for p in layer.parameters(): p.requires_grad = False return nn.Sequential(*features), classifier
def val(self): """ Validation function during the train phase. """ self.pose_net.eval() start_time = time.time() with torch.no_grad(): for i, data_dict in enumerate(self.val_loader): # Forward pass. out = self.pose_net(data_dict) # Compute the loss of the val batch. loss_dict = self.pose_loss(out) self.val_losses.update({key: loss.item() for key, loss in loss_dict.items()}, data_dict['img'].size(0)) # Update the vars of the val phase. self.batch_time.update(time.time() - start_time) start_time = time.time() self.runner_state['val_loss'] = self.val_losses.avg['loss'] RunnerHelper.save_net(self, self.pose_net, val_loss=self.val_losses.avg['loss']) # Print the log info & reset the states. Log.info( 'Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t' 'Loss {0}\n'.format(self.val_losses.info(), batch_time=self.batch_time)) self.batch_time.reset() self.val_losses.reset() self.pose_net.train()
def get_trainloader(self): if self.configer.get('dataset', default=None) == 'default_cpm': dataset = DefaultCPMDataset(root_dir=self.configer.get( 'data', 'data_dir'), dataset='train', aug_transform=self.aug_train_transform, img_transform=self.img_transform, configer=self.configer) elif self.configer.get('dataset', default=None) == 'default_openpose': dataset = DefaultOpenPoseDataset( root_dir=self.configer.get('data', 'data_dir'), dataset='train', aug_transform=self.aug_train_transform, img_transform=self.img_transform, configer=self.configer) else: Log.error('{} dataset is invalid.'.format( self.configer.get('dataset', default=None))) exit(1) trainloader = data.DataLoader( dataset, batch_size=self.configer.get('train', 'batch_size'), shuffle=True, num_workers=self.configer.get('data', 'workers'), pin_memory=True, drop_last=self.configer.get('data', 'drop_last'), collate_fn=lambda *args: collate(*args, trans_dict=self.configer.get( 'train', 'data_transformer'))) return trainloader
def __read_json_file(self, root_dir, dataset): img_list = list() label_list = list() with open(os.path.join(root_dir, dataset, 'label.json'), 'r') as file_stream: items = json.load(file_stream) for item in items: img_path = os.path.join(root_dir, dataset, item['image_path']) if not os.path.exists(img_path): Log.warn('Image Path: {} not exists.'.format(img_path)) continue img_list.append(img_path) label_list.append(item['label']) if dataset == 'train' and self.configer.get('data', 'include_val'): with open(os.path.join(root_dir, 'val', 'label.json'), 'r') as file_stream: items = json.load(file_stream) for item in items: img_path = os.path.join(root_dir, 'val', item['image_path']) if not os.path.exists(img_path): Log.warn('Image Path: {} not exists.'.format(img_path)) continue img_list.append(img_path) label_list.append(item['label']) return img_list, label_list
def get_trainloader(self): if self.configer.get('train.loader', default=None) in [None, 'default']: dataset = DefaultLoader(root_dir=self.configer.get( 'data', 'data_dir'), dataset='train', aug_transform=self.aug_train_transform, img_transform=self.img_transform, configer=self.configer) sampler = None if self.configer.get('network.distributed'): sampler = torch.utils.data.distributed.DistributedSampler( dataset) trainloader = data.DataLoader( dataset, sampler=sampler, batch_size=self.configer.get('train', 'batch_size'), shuffle=(sampler is None), num_workers=self.configer.get('data', 'workers'), pin_memory=True, drop_last=self.configer.get('data', 'drop_last'), collate_fn=lambda *args: collate( *args, trans_dict=self.configer.get('train', 'data_transformer'))) return trainloader else: Log.error('{} train loader is invalid.'.format( self.configer.get('train', 'loader'))) exit(1)
def init_weights(net, init_type='normal', init_gain=0.02): """Initialize network weights. Parameters: net (network) -- network to be initialized init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal init_gain (float) -- scaling factor for normal, xavier and orthogonal. We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might work better for some applications. Feel free to try yourself. """ def init_func(m): # define the initialization function classname = m.__class__.__name__ if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1): if init_type == 'normal': init.normal_(m.weight.data, 0.0, init_gain) elif init_type == 'xavier': init.xavier_normal_(m.weight.data, gain=init_gain) elif init_type == 'kaiming': init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') elif init_type == 'orthogonal': init.orthogonal_(m.weight.data, gain=init_gain) else: raise NotImplementedError( 'initialization method [%s] is not implemented' % init_type) if hasattr(m, 'bias') and m.bias is not None: init.constant_(m.bias.data, 0.0) elif classname.find( 'BatchNorm2d' ) != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies. init.normal_(m.weight.data, 1.0, init_gain) init.constant_(m.bias.data, 0.0) Log.info('initialize network with {}'.format(init_type)) net.apply(init_func) # apply the initialization function <init_func>
def get_valloader(self): if self.configer.get('dataset', default=None) in [None, 'default']: dataset = DefaultDataset(root_dir=self.configer.get( 'data', 'data_dir'), dataset='val', aug_transform=self.aug_val_transform, img_transform=self.img_transform, label_transform=self.label_transform, configer=self.configer) elif self.configer.get('dataset', default=None) == 'cityscapes': dataset = CityscapesDataset(root_dir=self.configer.get( 'data', 'data_dir'), dataset='val', aug_transform=self.aug_val_transform, img_transform=self.img_transform, label_transform=self.label_transform, configer=self.configer) else: Log.error('{} dataset is invalid.'.format( self.configer.get('dataset'))) exit(1) valloader = data.DataLoader( dataset, batch_size=self.configer.get('val', 'batch_size'), shuffle=False, num_workers=self.configer.get('data', 'workers'), pin_memory=True, collate_fn=lambda *args: collate(*args, trans_dict=self.configer.get( 'val', 'data_transformer'))) return valloader
def _make_parallel(runner, net): if runner.configer.get('network.distributed', default=False): #print('n1') from apex.parallel import DistributedDataParallel #print('n2') if runner.configer.get('network.syncbn', default=False): Log.info('Converting syncbn model...') from apex.parallel import convert_syncbn_model net = convert_syncbn_model(net) torch.cuda.set_device(runner.configer.get('local_rank')) torch.distributed.init_process_group(backend='nccl', init_method='env://') net = DistributedDataParallel(net.cuda(), delay_allreduce=True) return net net = net.to( torch.device( 'cpu' if runner.configer.get('gpu') is None else 'cuda')) if len(runner.configer.get('gpu')) > 1: from exts.tools.parallel.data_parallel import ParallelModel return ParallelModel(net, gather_=runner.configer.get( 'network', 'gather')) return net
def train(runner): Log.info('Training start...') if runner.configer.get('network', 'resume') is not None and runner.configer.get( 'network', 'resume_val'): runner.val() if runner.configer.get('solver', 'lr')['metric'] == 'epoch': while runner.runner_state['epoch'] < runner.configer.get( 'solver', 'max_epoch'): if runner.configer.get('network.distributed'): runner.train_loader.sampler.set_epoch( runner.runner_state['epoch']) runner.train() if runner.runner_state['epoch'] == runner.configer.get( 'solver', 'max_epoch'): runner.val() break else: while runner.runner_state['iters'] < runner.configer.get( 'solver', 'max_iters'): if runner.configer.get('network.distributed'): runner.train_loader.sampler.set_epoch( runner.runner_state['epoch']) runner.train() if runner.runner_state['iters'] == runner.configer.get( 'solver', 'max_iters'): runner.val() break Log.info('Training end...')
def vis_peaks(self, heatmap_in, ori_img_in, name='default', sub_dir='peaks'): base_dir = os.path.join(self.configer.get('project_dir'), POSE_DIR, sub_dir) if not os.path.exists(base_dir): Log.error('Dir:{} not exists!'.format(base_dir)) os.makedirs(base_dir) if not isinstance(heatmap_in, np.ndarray): if len(heatmap_in.size()) != 3: Log.error('Heatmap size is not valid.') exit(1) heatmap = heatmap_in.clone().data.cpu().numpy().transpose(1, 2, 0) else: heatmap = heatmap_in.copy() if not isinstance(ori_img_in, np.ndarray): ori_img = DeNormalize(div_value=self.configer.get('normalize', 'div_value'), mean=self.configer.get('normalize', 'mean'), std=self.configer.get('normalize', 'std'))(ori_img_in.clone()) ori_img = ori_img.data.cpu().squeeze().numpy().transpose(1, 2, 0).astype(np.uint8) ori_img = cv2.cvtColor(ori_img, cv2.COLOR_RGB2BGR) else: ori_img = ori_img_in.copy() for j in range(self.configer.get('data', 'num_kpts')): peaks = self.__get_peaks(heatmap[:, :, j]) for peak in peaks: ori_img = cv2.circle(ori_img, (peak[0], peak[1]), self.configer.get('vis', 'circle_radius'), self.configer.get('details', 'color_list')[j], thickness=-1) cv2.imwrite(os.path.join(base_dir, '{}_{}.jpg'.format(name, j)), ori_img)
def get_trainloader(self): if self.configer.get('train.loader', default=None) in [None, 'default']: trainloader = data.DataLoader( DefaultLoader(root_dir=self.configer.get('data', 'data_dir'), dataset='train', aug_transform=self.aug_train_transform, img_transform=self.img_transform, configer=self.configer), batch_size=self.configer.get('train', 'batch_size'), shuffle=True, num_workers=self.configer.get('data', 'workers'), pin_memory=True, drop_last=self.configer.get('data', 'drop_last'), collate_fn=lambda *args: collate( *args, trans_dict=self.configer.get('train', 'data_transformer') ) ) return trainloader elif self.configer.get('train', 'loader') == 'fasterrcnn': trainloader = data.DataLoader( FasterRCNNLoader(root_dir=self.configer.get('data', 'data_dir'), dataset='train', aug_transform=self.aug_train_transform, img_transform=self.img_transform, configer=self.configer), batch_size=self.configer.get('train', 'batch_size'), shuffle=True, num_workers=self.configer.get('data', 'workers'), pin_memory=True, drop_last=self.configer.get('data', 'drop_last'), collate_fn=lambda *args: collate( *args, trans_dict=self.configer.get('train', 'data_transformer') ) ) return trainloader else: Log.error('{} train loader is invalid.'.format(self.configer.get('train', 'loader'))) exit(1)
def cv2_read_image(image_path, mode='RGB'): if ImageHelper.is_zip_path(image_path): if mode == 'RGB': return ImageHelper.bgr2rgb(ZipReader.imread(image_path, mode)) elif mode == 'BGR': return ZipReader.imread(image_path, mode) elif mode == 'P': return ZipReader.imread(image_path, mode) else: Log.error('Not support mode {}'.format(mode)) exit(1) else: img_bgr = cv2.imread(image_path, cv2.IMREAD_COLOR) if mode == 'RGB': return ImageHelper.bgr2rgb(img_bgr) elif mode == 'BGR': return img_bgr elif mode == 'P': return ImageHelper.to_np(Image.open(image_path).convert('P')) else: Log.error('Not support mode {}'.format(mode)) exit(1)
def vis_bboxes(self, image_in, bboxes_list, name='default', sub_dir='bbox'): """ Show the diff bbox of individuals. """ base_dir = os.path.join(self.configer.get('project_dir'), DET_DIR, sub_dir) if isinstance(image_in, Image.Image): image = ImageHelper.rgb2bgr(ImageHelper.to_np(image_in)) else: image = image_in.copy() if not os.path.exists(base_dir): log.error('Dir:{} not exists!'.format(base_dir)) os.makedirs(base_dir) img_path = os.path.join( base_dir, name if ImageHelper.is_img(name) else '{}.jpg'.format(name)) for bbox in bboxes_list: image = cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2) cv2.imwrite(img_path, image)
def __test_img(self, image_path, save_path): Log.info('Image Path: {}'.format(image_path)) ori_image = ImageHelper.read_image(image_path, tool=self.configer.get('data', 'image_tool'), mode=self.configer.get('data', 'input_mode')) ori_width, ori_height = ImageHelper.get_size(ori_image) ori_img_bgr = ImageHelper.get_cv2_bgr(ori_image, mode=self.configer.get('data', 'input_mode')) heatmap_avg = np.zeros((ori_height, ori_width, self.configer.get('network', 'heatmap_out'))) for i, scale in enumerate(self.configer.get('test', 'scale_search')): image = self.blob_helper.make_input(ori_image, input_size=self.configer.get('test', 'input_size'), scale=scale) with torch.no_grad(): heatmap_out_list = self.pose_net(image) heatmap_out = heatmap_out_list[-1] # extract outputs, resize, and remove padding heatmap = heatmap_out.squeeze(0).cpu().numpy().transpose(1, 2, 0) heatmap = cv2.resize(heatmap, (ori_width, ori_height), interpolation=cv2.INTER_CUBIC) heatmap_avg = heatmap_avg + heatmap / len(self.configer.get('test', 'scale_search')) all_peaks = self.__extract_heatmap_info(heatmap_avg) image_canvas = self.__draw_key_point(all_peaks, ori_img_bgr) ImageHelper.save(image_canvas, save_path)
def get_valloader(self, dataset=None): dataset = 'val' if dataset is None else dataset if self.configer.get('val.loader', default=None) in [None, 'default']: valloader = data.DataLoader( DefaultLoader(root_dir=self.configer.get('data', 'data_dir'), dataset=dataset, aug_transform=self.aug_val_transform, img_transform=self.img_transform, configer=self.configer), batch_size=self.configer.get('val', 'batch_size'), shuffle=False, num_workers=self.configer.get('data', 'workers'), pin_memory=True, collate_fn=lambda *args: collate( *args, trans_dict=self.configer.get('val', 'data_transformer') ) ) return valloader elif self.configer.get('val', 'loader') == 'fasterrcnn': valloader = data.DataLoader( FasterRCNNLoader(root_dir=self.configer.get('data', 'data_dir'), dataset=dataset, aug_transform=self.aug_val_transform, img_transform=self.img_transform, configer=self.configer), batch_size=self.configer.get('val', 'batch_size'), shuffle=False, num_workers=self.configer.get('data', 'workers'), pin_memory=True, collate_fn=lambda *args: collate( *args, trans_dict=self.configer.get('val', 'data_transformer') ) ) return valloader else: Log.error('{} val loader is invalid.'.format(self.configer.get('val', 'loader'))) exit(1)
def check_and_load(path: str, remote_url: str = None) -> None: """ Check if the file at the path exists. If not, and remote_url is not None, then offer to load it. """ if not os.path.exists(path): print("\nError: File {} does not exist.".format(path)) if remote_url is not None: while True: response = input( "Do you want to load it from {}? [y/n]: ".format( remote_url)) if response == "y": print( "Loading, please wait! (The file is pretty big so this could take a while...)" ) pather.create(path) urllib.request.urlretrieve(remote_url, path) Logger.log_field("Successfully Loaded", path) if response == "n": break print("Response is not valid. Please enter y or n.\n") raise Exception("File not found, unable to proceed.")
def train(self): """ Train function of every epoch during train phase. """ self.det_net.train() start_time = time.time() # Adjust the learning rate after every epoch. self.runner_state['epoch'] += 1 # data_tuple: (inputs, heatmap, maskmap, vecmap) for i, data_dict in enumerate(self.train_loader): Trainer.update(self, warm_list=(0, ), warm_lr_list=(self.configer.get('solver', 'lr')['base_lr'], ), solver_dict=self.configer.get('solver')) self.data_time.update(time.time() - start_time) # Forward pass. out_dict = self.det_net(data_dict) # Compute the loss of the train batch & backward. loss = out_dict['loss'].mean() self.train_losses.update(loss.item(), len(DCHelper.tolist(data_dict['meta']))) self.optimizer.zero_grad() loss.backward() self.optimizer.step() # Update the vars of the train phase. self.batch_time.update(time.time() - start_time) start_time = time.time() self.runner_state['iters'] += 1 # Print the log info & reset the states. if self.runner_state['iters'] % self.configer.get( 'solver', 'display_iter') == 0: Log.info( 'Train Epoch: {0}\tTrain Iteration: {1}\t' 'Time {batch_time.sum:.3f}s / {2}iters, ({batch_time.avg:.3f})\t' 'Data load {data_time.sum:.3f}s / {2}iters, ({data_time.avg:3f})\n' 'Learning rate = {3}\tLoss = {loss.val:.8f} (ave = {loss.avg:.8f})\n' .format(self.runner_state['epoch'], self.runner_state['iters'], self.configer.get('solver', 'display_iter'), RunnerHelper.get_lr(self.optimizer), batch_time=self.batch_time, data_time=self.data_time, loss=self.train_losses)) self.batch_time.reset() self.data_time.reset() self.train_losses.reset() if self.configer.get('solver', 'lr')['metric'] == 'iters' \ and self.runner_state['iters'] == self.configer.get('solver', 'max_iters'): break # Check to val the current model. if self.runner_state['iters'] % self.configer.get( 'solver', 'test_interval') == 0: self.val()
def _make_parallel(runner, net): if runner.configer.get('network.distributed', default=False): local_rank = runner.configer.get('local_rank') torch.cuda.set_device(local_rank) torch.distributed.init_process_group(backend='nccl', init_method='env://') if runner.configer.get('network.syncbn', default=False): Log.info('Converting syncbn model...') net = nn.SyncBatchNorm.convert_sync_batchnorm(net) net = nn.parallel.DistributedDataParallel( net.cuda(), find_unused_parameters=True, device_ids=[local_rank], output_device=local_rank) # if runner.configer.get('network.syncbn', default=False): # Log.info('Converting syncbn model...') # from apex.parallel import convert_syncbn_model # net = convert_syncbn_model(net) # from apex.parallel import DistributedDataParallel # net = DistributedDataParallel(net.cuda(), delay_allreduce=True) return net net = net.to( torch.device( 'cpu' if runner.configer.get('gpu') is None else 'cuda')) from exts.tools.parallel.data_parallel import ParallelModel return ParallelModel(net, gather_=runner.configer.get('network', 'gather'))
def get_valloader(self, dataset=None): dataset = 'val' if dataset is None else dataset if self.configer.get('val.loader', default=None) in [None, 'default']: dataset = DefaultLoader(root_dir=self.configer.get( 'data', 'data_dir'), dataset=dataset, aug_transform=self.aug_val_transform, img_transform=self.img_transform, configer=self.configer) sampler = None if self.configer.get('network.distributed'): sampler = torch.utils.data.distributed.DistributedSampler( dataset) valloader = data.DataLoader( dataset, sampler=sampler, batch_size=self.configer.get('val', 'batch_size'), shuffle=False, num_workers=self.configer.get('data', 'workers'), pin_memory=True, collate_fn=lambda *args: collate( *args, trans_dict=self.configer.get('val', 'data_transformer'))) return valloader else: Log.error('{} val loader is invalid.'.format( self.configer.get('val', 'loader'))) exit(1)
def val(self): """ Validation function during the train phase. """ self.gan_net.eval() start_time = time.time() for j, data_dict in enumerate(self.val_loader): with torch.no_grad(): # Forward pass. out_dict = self.gan_net(data_dict) # Compute the loss of the val batch. self.val_losses.update( out_dict['loss_G'].mean().item() + out_dict['loss_D'].mean().item(), len(DCHelper.tolist(data_dict['meta']))) # Update the vars of the val phase. self.batch_time.update(time.time() - start_time) start_time = time.time() RunnerHelper.save_net(self, self.gan_net, val_loss=self.val_losses.avg) # Print the log info & reset the states. Log.info('Test Time {batch_time.sum:.3f}s, ({batch_time.avg:.3f})\t' 'Loss {loss.avg:.8f}\n'.format(batch_time=self.batch_time, loss=self.val_losses)) self.batch_time.reset() self.val_losses.reset() self.gan_net.train()
def __init__(self, configer): self.configer = configer if self.configer.get('data', 'image_tool') == 'pil': self.aug_train_transform = pil_aug_trans.PILAugCompose( self.configer, split='train') elif self.configer.get('data', 'image_tool') == 'cv2': self.aug_train_transform = cv2_aug_trans.CV2AugCompose( self.configer, split='train') else: Log.error('Not support {} image tool.'.format( self.configer.get('data', 'image_tool'))) exit(1) if self.configer.get('data', 'image_tool') == 'pil': self.aug_val_transform = pil_aug_trans.PILAugCompose(self.configer, split='val') elif self.configer.get('data', 'image_tool') == 'cv2': self.aug_val_transform = cv2_aug_trans.CV2AugCompose(self.configer, split='val') else: Log.error('Not support {} image tool.'.format( self.configer.get('data', 'image_tool'))) exit(1) self.img_transform = trans.Compose([ trans.ToTensor(), trans.Normalize(**self.configer.get('data', 'normalize')), ])
def get_backbone(self, **params): backbone = self.configer.get('network', 'backbone') model = None if 'vgg' in backbone: model = VGGBackbone(self.configer)(**params) elif 'darknet' in backbone: model = DarkNetBackbone(self.configer)(**params) elif 'resnet' in backbone: model = ResNetBackbone(self.configer)(**params) elif 'mobilenet' in backbone: model = MobileNetBackbone(self.configer)(*params) elif 'densenet' in backbone: model = DenseNetBackbone(self.configer)(**params) elif 'squeezenet' in backbone: model = SqueezeNetBackbone(self.configer)(**params) elif 'shufflenet' in backbone: model = ShuffleNetv2Backbone(self.configer)(**params) else: Log.error('Backbone {} is invalid.'.format(backbone)) exit(1) return model
def get_valloader(self, dataset=None): dataset = 'val' if dataset is None else dataset if self.configer.get('dataset', default=None) == 'default_cpm': dataset = DefaultDataset(root_dir=self.configer.get( 'data', 'data_dir'), dataset=dataset, aug_transform=self.aug_val_transform, img_transform=self.img_transform, configer=self.configer) elif self.configer.get('dataset', default=None) == 'default_openpose': dataset = DefaultOpenPoseDataset( root_dir=self.configer.get('data', 'data_dir'), dataset=dataset, aug_transform=self.aug_val_transform, img_transform=self.img_transform, configer=self.configer), else: Log.error('{} dataset is invalid.'.format( self.configer.get('dataset'))) exit(1) valloader = data.DataLoader( dataset, batch_size=self.configer.get('val', 'batch_size'), shuffle=False, num_workers=self.configer.get('data', 'workers'), pin_memory=True, collate_fn=lambda *args: collate(*args, trans_dict=self.configer.get( 'val', 'data_transformer'))) return valloader
def __list_dirs(self, root_dir, dataset): imgA_list = list() imgB_list = list() imageA_dir = os.path.join(root_dir, dataset, 'imageA') imageB_dir = os.path.join(root_dir, dataset, 'imageB') for file_name in os.listdir(imageA_dir): image_name = '.'.join(file_name.split('.')[:-1]) imgA_path = ImageHelper.imgpath(imageA_dir, image_name) imgB_path = ImageHelper.imgpath(imageB_dir, image_name) if not os.path.exists(imgA_path) or not os.path.exists(imgB_path): Log.warn('Img Path: {} not exists.'.format(imgA_path)) continue imgA_list.append(imgA_path) imgB_list.append(imgB_path) if dataset == 'train' and self.configer.get('data', 'include_val'): imageA_dir = os.path.join(root_dir, 'val/imageA') imageB_dir = os.path.join(root_dir, 'val/imageB') for file_name in os.listdir(imageA_dir): image_name = '.'.join(file_name.split('.')[:-1]) imgA_path = ImageHelper.imgpath(imageA_dir, image_name) imgB_path = ImageHelper.imgpath(imageB_dir, image_name) if not os.path.exists(imgA_path) or not os.path.exists( imgB_path): Log.warn('Img Path: {} not exists.'.format(imgA_path)) continue imgA_list.append(imgA_path) imgB_list.append(imgB_path) return imgA_list, imgB_list
def load_net(runner, net, model_path=None, map_location='cpu'): if model_path is not None or runner.configer.get('network', 'resume') is not None: resume_path = runner.configer.get('network', 'resume') resume_path = model_path if model_path is not None else resume_path Log.info('Resuming from {}'.format(resume_path)) resume_dict = torch.load(resume_path, map_location=map_location) if 'state_dict' in resume_dict: checkpoint_dict = resume_dict['state_dict'] elif 'model' in resume_dict: checkpoint_dict = resume_dict['model'] elif isinstance(resume_dict, OrderedDict): checkpoint_dict = resume_dict else: raise RuntimeError( 'No state_dict found in checkpoint file {}'.format(runner.configer.get('network', 'resume'))) # load state_dict if hasattr(net, 'module'): RunnerHelper.load_state_dict(net.module, checkpoint_dict, runner.configer.get('network', 'resume_strict')) else: RunnerHelper.load_state_dict(net, checkpoint_dict, runner.configer.get('network', 'resume_strict')) if runner.configer.get('network', 'resume_continue'): # runner.configer.resume(resume_dict['config_dict']) runner.runner_state = resume_dict['runner_state'] net = RunnerHelper._make_parallel(runner, net) return net
def vis_default_bboxes(self, ori_img_in, default_bboxes, labels, name='default', sub_dir='encode'): base_dir = os.path.join(self.configer.get('project_dir'), DET_DIR, sub_dir) if not os.path.exists(base_dir): log.error('Dir:{} not exists!'.format(base_dir)) os.makedirs(base_dir) if not isinstance(ori_img_in, np.ndarray): ori_img = DeNormalize( div_value=self.configer.get('normalize', 'div_value'), mean=self.configer.get('normalize', 'mean'), std=self.configer.get('normalize', 'std'))(ori_img_in.clone()) ori_img = ori_img.data.cpu().squeeze().numpy().transpose( 1, 2, 0).astype(np.uint8) ori_img = cv2.cvtColor(ori_img, cv2.COLOR_RGB2BGR) else: ori_img = ori_img_in.copy() assert labels.size(0) == default_bboxes.size(0) bboxes = torch.cat([ default_bboxes[:, :2] - default_bboxes[:, 2:] / 2, default_bboxes[:, :2] + default_bboxes[:, 2:] / 2 ], 1) height, width, _ = ori_img.shape for i in range(labels.size(0)): if labels[i] == 0: continue class_name = self.configer.get('details', 'name_seq')[labels[i] - 1] color_num = len(self.configer.get('details', 'color_list')) cv2.rectangle( ori_img, (int(bboxes[i][0] * width), int(bboxes[i][1] * height)), (int(bboxes[i][2] * width), int(bboxes[i][3] * height)), color=self.configer.get( 'details', 'color_list')[(labels[i] - 1) % color_num], thickness=3) cv2.putText(ori_img, class_name, (int(bboxes[i][0] * width) + 5, int(bboxes[i][3] * height) - 5), cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=self.configer.get('details', 'color_list')[(labels[i] - 1) % color_num], thickness=2) img_path = os.path.join(base_dir, '{}.jpg'.format(name)) cv2.imwrite(img_path, ori_img)
def test(self, test_dir, out_dir): for _, data_dict in enumerate( self.test_loader.get_testloader(test_dir=test_dir)): data_dict['testing'] = True data_dict = RunnerHelper.to_device(self, data_dict) out_dict = self.det_net(data_dict) meta_list = DCHelper.tolist(data_dict['meta']) test_indices_and_rois, test_roi_locs, test_roi_scores, test_rois_num = out_dict[ 'test_group'] batch_detections = self.decode(test_roi_locs, test_roi_scores, test_indices_and_rois, test_rois_num, self.configer, meta_list) for i in range(len(meta_list)): ori_img_bgr = ImageHelper.read_image(meta_list[i]['img_path'], tool='cv2', mode='BGR') json_dict = self.__get_info_tree(batch_detections[i]) image_canvas = self.det_parser.draw_bboxes( ori_img_bgr.copy(), json_dict, conf_threshold=self.configer.get('res', 'vis_conf_thre')) ImageHelper.save(image_canvas, save_path=os.path.join( out_dir, 'vis/{}.png'.format( meta_list[i]['filename']))) Log.info('Json Path: {}'.format( os.path.join( out_dir, 'json/{}.json'.format(meta_list[i]['filename'])))) JsonHelper.save_file(json_dict, save_path=os.path.join( out_dir, 'json/{}.json'.format( meta_list[i]['filename'])))
def _process_slide(self, slide: Slide): slide_path = os.path.join(self.output_path, slide.name) slide.slide_path = slide_path os.mkdir(slide_path) # Draw the slide image. image = slide.image image_path = os.path.join(slide_path, "image.png") Logger.log_special("Scanning {}".format(slide.name), with_gap=True) # Predict the cell masks from an image. predict_image = image if True: equalizer_image = self.equalizer.create_equalized_image(image) predict_image = equalizer_image # Get the sample prediction. prediction = self.net.cycle_predict(predict_image, None) slide.cells = self._process_prediction(slide, slide_path, prediction) self._draw_prediction_mask(image, slide_path, prediction) pather.create("output/summary") self.reporter.produce(slide, "output/summary") cv2.imwrite(image_path, equalizer_image)