def get_pred(batch_data): # process data batch_data = batch_data[0] segSize = (batch_data['img_ori'].shape[0], batch_data['img_ori'].shape[1]) img_resized_list = batch_data['img_data'] with torch.no_grad(): scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0], segSize[1]) scores = async_copy_to(scores, gpu) for img in img_resized_list: feed_dict = batch_data.copy() feed_dict['img_data'] = img del feed_dict['img_ori'] del feed_dict['info'] feed_dict = async_copy_to(feed_dict, gpu) # forward pass pred_tmp = segmentation_module(feed_dict, segSize=segSize) scores = scores + pred_tmp / len(cfg.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) writer.write(pred) pred = as_numpy(pred.squeeze(0).cpu()) pbar.update(1) return pred
def test(segmentation_module, loader, gpu): segmentation_module.eval() pbar = tqdm(total=len(loader)) for batch_data in loader: # process data batch_data = batch_data[0] segSize = (batch_data['img_ori'].shape[0], batch_data['img_ori'].shape[1]) img_resized_list = batch_data['img_data'] with torch.no_grad(): scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0], segSize[1]) scores = async_copy_to(scores, gpu) for img in img_resized_list: feed_dict = batch_data.copy() feed_dict['img_data'] = img del feed_dict['img_ori'] del feed_dict['info'] feed_dict = async_copy_to(feed_dict, gpu) # forward pass pred_tmp = segmentation_module(feed_dict, segSize=segSize) scores = scores + pred_tmp / len(cfg.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) pred = as_numpy(pred.squeeze(0).cpu()) # visualization visualize_result((batch_data['img_ori'], batch_data['info']), pred, cfg) pbar.update(1)
def huawei_seg(imgs, segmentation_module): # cfg.merge_from_file("/home/mind/model/config/ade20k-hrnetv2-huawei.yaml") imgs = [imgs] cfg.list_test = [{'fpath_img': x} for x in imgs] # Dataset and Loader dataset_test = InferDataset(cfg.list_test, cfg.DATASET) loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() loader = loader_test # Main loop segmentation_module.eval() pbar = tqdm(total=len(loader)) for batch_data in loader: # process data batch_data = batch_data[0] segSize = (batch_data['img_ori'].shape[0], batch_data['img_ori'].shape[1]) img_resized_list = batch_data['img_data'] with torch.no_grad(): scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0], segSize[1]) scores = async_copy_to(scores, 0) for img in img_resized_list: feed_dict = batch_data.copy() feed_dict['img_data'] = img del feed_dict['img_ori'] del feed_dict['info'] feed_dict = async_copy_to(feed_dict, 0) # forward pass pred_tmp = segmentation_module(feed_dict, segSize=segSize) # print(pred_tmp.shape)#torch.Size([1, 2, 1024, 1024]) scores = scores + pred_tmp / len(cfg.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) pred = as_numpy(pred.squeeze(0).cpu()) # # visualize_result( # (batch_data['img_ori'], batch_data['info']), # pred, # cfg # ) pbar.update(1) # return pred
def test(segmentation_module, loader, device, names): segmentation_module.eval() pbar = tqdm(total=len(loader)) re = [] for batch_data in loader: # process data batch_data = batch_data[0] segSize = (batch_data['img_ori'].shape[0], batch_data['img_ori'].shape[1]) img_resized_list = batch_data['img_data'] with torch.no_grad(): scores = torch.zeros(1, cfg_ss.DATASET.num_class, segSize[0], segSize[1]) if torch.cuda.is_available(): scores = async_copy_to(scores, device) for img in img_resized_list: feed_dict = batch_data.copy() feed_dict['img_data'] = img del feed_dict['img_ori'] del feed_dict['info'] if torch.cuda.is_available(): feed_dict = async_copy_to(feed_dict, device) # forward pass pred_tmp = segmentation_module(feed_dict, segSize=segSize) scores = scores + pred_tmp / len(cfg_ss.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) pred = as_numpy(pred.squeeze(0).cpu()) pred = np.int32(pred) pixs = pred.size uniques, counts = np.unique(pred, return_counts=True) pred_ratios = {} for idx in np.argsort(counts)[::-1]: name = names[uniques[idx] + 1] ratio = counts[idx] / pixs * 100 if ratio > 0.1: pred_ratios[name] = ratio re.append({ "original_img": batch_data['img_ori'], "img_inf": batch_data['info'], "pred_result": pred, "pred_ratio": pred_ratios }) pbar.update(1) return re
def run_inference_for_single_image(self, image): preproc_data = preprocess_image(image) batch_data = preproc_data seg_size = (batch_data['img_ori'].shape[0], batch_data['img_ori'].shape[1]) img_resized_list = batch_data['img_data'] if self.gpu is not None: with torch.no_grad(): # scores = torch.zeros(1, cfg.DATASET.num_class, seg_size[0], seg_size[1]) scores = torch.cuda.FloatTensor(1, cfg.DATASET.num_class, seg_size[0], seg_size[1]).fill_(0) scores = async_copy_to(scores, self.gpu) for img in img_resized_list: feed_dict = {'img_data': img} feed_dict = async_copy_to(feed_dict, self.gpu) # forward pass pred_tmp = self.model(feed_dict, segSize=seg_size) scores = scores + pred_tmp / len(cfg.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) pred = pred.squeeze(0).cpu().numpy() else: scores = torch.zeros(1, cfg.DATASET.num_class, seg_size[0], seg_size[1]) for img in img_resized_list: feed_dict = {'img_data': img} # forward pass pred_tmp = self.model(feed_dict, segSize=seg_size) scores = scores + pred_tmp / len(cfg.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) pred = pred.squeeze(0).numpy() return pred
def evaluate(segmentation_module, loader, cfg, gpu_id, result_queue): segmentation_module.eval() for index, batch_data in enumerate(loader): # process data batch_data = batch_data[0] seg_label = as_numpy(batch_data['seg_label'][0]) img_resized_list = batch_data['img_data'] with torch.no_grad(): segSize = (seg_label.shape[0], seg_label.shape[1]) scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0], segSize[1]) scores = async_copy_to(scores, gpu_id) for img in img_resized_list: feed_dict = batch_data.copy() feed_dict['img_data'] = img del feed_dict['img_ori'] del feed_dict['info'] feed_dict = async_copy_to(feed_dict, gpu_id) # forward pass scores_tmp = segmentation_module(feed_dict, segSize=segSize) scores = scores + scores_tmp / len(cfg.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) pred = as_numpy(pred.squeeze(0).cpu()) # calculate accuracy and SEND THEM TO MASTER acc, pix = accuracy(pred, seg_label) intersection, union = intersectionAndUnion(pred, seg_label, cfg.DATASET.num_class) result_queue.put_nowait((acc, pix, intersection, union)) # visualization if cfg.VAL.visualize: visualize_result( (batch_data['img_ori'], seg_label, batch_data['info']), pred, os.path.join(cfg.DIR, 'result'))
def huawei_seg(imgs): parser = argparse.ArgumentParser( description="PyTorch Semantic Segmentation Testing") parser.add_argument( "--cfg", default="config/ade20k-hrnetv2-huawei.yaml", metavar="FILE", help="path to config file", type=str, ) parser.add_argument("--gpu", default=0, type=int, help="gpu id for evaluation") args = parser.parse_args() cfg.merge_from_file(args.cfg) cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower() cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower() # absolute paths of model weights cfg.MODEL.weights_encoder = os.path.join(cfg.DIR, 'encoder_' + cfg.TEST.checkpoint) cfg.MODEL.weights_decoder = os.path.join(cfg.DIR, 'decoder_' + cfg.TEST.checkpoint) # imgs = [imgs] cfg.list_test = [{'fpath_img': x} for x in imgs] torch.cuda.set_device(args.gpu) # Network Builders net_encoder = ModelBuilder.build_encoder(arch=cfg.MODEL.arch_encoder, fc_dim=cfg.MODEL.fc_dim, weights=cfg.MODEL.weights_encoder) net_decoder = ModelBuilder.build_decoder(arch=cfg.MODEL.arch_decoder, fc_dim=cfg.MODEL.fc_dim, num_class=cfg.DATASET.num_class, weights=cfg.MODEL.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) # Dataset and Loader dataset_test = InferDataset(cfg.list_test, cfg.DATASET) loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() loader = loader_test # Main loop segmentation_module.eval() pbar = tqdm(total=len(loader)) for batch_data in loader: # process data batch_data = batch_data[0] segSize = (batch_data['img_ori'].shape[0], batch_data['img_ori'].shape[1]) img_resized_list = batch_data['img_data'] with torch.no_grad(): scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0], segSize[1]) scores = async_copy_to(scores, args.gpu) for img in img_resized_list: feed_dict = batch_data.copy() feed_dict['img_data'] = img del feed_dict['img_ori'] del feed_dict['info'] feed_dict = async_copy_to(feed_dict, args.gpu) # forward pass pred_tmp = segmentation_module(feed_dict, segSize=segSize) # print(pred_tmp.shape)#torch.Size([1, 2, 1024, 1024]) scores = scores + pred_tmp / len(cfg.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) pred = as_numpy(pred.squeeze(0).cpu()) # visualization visualize_result((batch_data['img_ori'], batch_data['info']), pred, cfg) pbar.update(1) # return pred
def evaluate(segmentation_module, loader, cfg, gpu): acc_meter = AverageMeter() intersection_meter = AverageMeter() union_meter = AverageMeter() time_meter = AverageMeter() segmentation_module.eval() pbar = tqdm(total=len(loader)) for batch_data in loader: # process data batch_data = batch_data[0] seg_label = as_numpy(batch_data['seg_label'][0]) img_resized_list = batch_data['img_data'] torch.cuda.synchronize() tic = time.perf_counter() with torch.no_grad(): segSize = (seg_label.shape[0], seg_label.shape[1]) scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0], segSize[1]) scores = async_copy_to(scores, gpu) for img in img_resized_list: feed_dict = batch_data.copy() feed_dict['img_data'] = img del feed_dict['img_ori'] del feed_dict['info'] feed_dict = async_copy_to(feed_dict, gpu) # forward pass scores_tmp = segmentation_module(feed_dict, segSize=segSize) scores = scores + scores_tmp / len(cfg.DATASET.imgSizes) _, pred = torch.max(scores, dim=1) pred = as_numpy(pred.squeeze(0).cpu()) torch.cuda.synchronize() time_meter.update(time.perf_counter() - tic) # calculate accuracy acc, pix = accuracy(pred, seg_label) intersection, union = intersectionAndUnion(pred, seg_label, cfg.DATASET.num_class) acc_meter.update(acc, pix) intersection_meter.update(intersection) union_meter.update(union) # visualization if cfg.VAL.visualize: visualize_result( (batch_data['img_ori'], seg_label, batch_data['info']), pred, os.path.join(cfg.DIR, 'result') ) pbar.update(1) # summary iou = intersection_meter.sum / (union_meter.sum + 1e-10) for i, _iou in enumerate(iou): print('class [{}], IoU: {:.4f}'.format(i, _iou)) print('[Eval Summary]:') print('Mean IoU: {:.4f}, Accuracy: {:.2f}%, Inference Time: {:.4f}s' .format(iou.mean(), acc_meter.average()*100, time_meter.average()))