def main(): parser = ArgumentParser() parser.add_argument('img', help='Image file') parser.add_argument('config', help='Config file') parser.add_argument('checkpoint', help='Checkpoint file') parser.add_argument( '--device', default='cuda:0', help='Device used for inference') parser.add_argument( '--palette', default='cityscapes', help='Color palette used for segmentation map') parser.add_argument( '--opacity', type=float, default=0.5, help='Opacity of painted segmentation map. In (0, 1] range.') args = parser.parse_args() # build the model from a config file and a checkpoint file model = init_segmentor(args.config, args.checkpoint, device=args.device) # test a single image result = inference_segmentor(model, args.img) # show the results show_result_pyplot( model, args.img, result, get_palette(args.palette), opacity=args.opacity)
def main(): parser = ArgumentParser() parser.add_argument('img', help='Image file') parser.add_argument('config', help='Config file') parser.add_argument('checkpoint', help='Checkpoint file') parser.add_argument( '--device', default='cuda:0', help='Device used for inference') parser.add_argument( '--palette', default='ade', help='Color palette used for segmentation map') args = parser.parse_args() # build the model from a config file and a checkpoint file model = init_segmentor(args.config, args.checkpoint, device=args.device) # test a single image demo_im_names = os.listdir(args.img) random.shuffle(demo_im_names) for im_name in demo_im_names: if 'png' in im_name or 'jpg' in im_name: full_name = os.path.join(args.img, im_name) result = inference_segmentor(model, full_name) # show the results pl=[[220, 220, 220],[17, 142, 35], [152, 251, 152], [0, 60, 100], [70, 130, 180], [220, 20, 20]] show_result_pyplot(model, full_name, result,pl)
def main(): parser = ArgumentParser() parser.add_argument('--img', help='Image file', default=img_path) parser.add_argument('--config', help='Config file', default=config) parser.add_argument('--checkpoint', help='Checkpoint file', default=ckpt) parser.add_argument( '--device', default='cuda:1', help='Device used for inference') parser.add_argument( '--palette', default=None, help='Color palette used for segmentation map') args = parser.parse_args() # build the model from a config file and a checkpoint file model = init_segmentor(args.config, args.checkpoint, device=args.device) if args.img=='': list_img=get_list_file_in_folder(img_dir) list_img=sorted(list_img) for img_ in list_img: img=os.path.join(img_dir,img_) print(img) result = inference_segmentor(model, img) show_result_pyplot(model, img, result, get_palette(args.palette)) else: result = inference_segmentor(model, args.img) show_result_pyplot(model, args.img, result, get_palette(args.palette))
def __init__(self): self.root_split_path = "../LidarDetector/data/kitti/training/" self.save_path = "../LidarDetector/data/kitti/training/painted_lidar_mm/" config_file = 'mmseg/configs/deeplabv3plus/deeplabv3plus_r101-d8_512x1024_80k_cityscapes.py' checkpoint_file = 'mmseg/checkpoints/deeplabv3plus_r101-d8_512x1024_80k_cityscapes_20200606_114143-068fcfe9.pth' self.model = init_segmentor(config_file, checkpoint_file, device='cuda:1')
def __init__(self, model_name, model_path): self.model_name = model_name self.model_path = model_path config_file = '/home/mind/model/deeplabv3plus512x512.py' self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") self.model = init_segmentor(config_file, self.model_path, device=self.device)
def init_model(transform): # use this format SegFormr.b1-512-ade in -t parameter to select a certain config/checkpoint (version, resolution, datasetType) = transform.split('-') config = "local_configs/segformer/" + version.upper( ) + "/segformer." + version + "." + resolution + "x" + resolution + "." + datasetType + ".160k.py" checkpoint = "checkpoints/segformer." + version + "." + resolution + "x" + resolution + "." + datasetType + ".160k.pth" # build the model from a config file and a checkpoint file model = init_segmentor(config, checkpoint, device='cuda:0') return (model), None
def main(): # # 定义 config 文件路径 # config_file = 'mmsegmentation/configs/deeplabv3plus/deeplabv3plus_r50-d8_512x1024_40k_cityscapes.py' # checkpoint_file = '../mmsegmentation/work_dirs/deeplabv3plus_r50-d8_512x1024_40k_cityscapes/latest.pth' # # # build the model from a config file and a checkpoint file # model = init_segmentor(config_file, checkpoint_file) # # # test a single image and show the results # img_path = 'demo.png' # or img = mmcv.imread(img), which will only load it once # img=cv2.imread(img_path) # print('img_cv.shape:',img.shape) # # imgPil=Image.open(img_path) # # imgPil=np.array(imgPil) # # print('imgPil.shape:',imgPil.shape) # # img_pil2cv = cv2.cvtColor(imgPil, cv2.COLOR_RGB2BGR) # # print('img_pil2cv.shape:',img_pil2cv.shape) # result = inference_segmentor(model, img_path) # print(result[0]) # #result=result[0] # #result=result.astype(np.int8) # #print(type(result[0][0])) # #print(result) # # visualize the results in a new window # #model.show_result(img, result, show=True) # # or save the visualization results to image files # #model.show_result(img, result, out_file='resultimg_pil2cv.jpg') config_file = 'config_huawei/deeplabv3plus512x512.py' checkpoint_file = '/home/admins/qyl/huawei_compete/mmsegmentation/work_dirs/deeplabv3plus_r50-d8_512x1024_40k_cityscapes/model_best.pth' # build the model from a config file and a checkpoint file model = init_segmentor(config_file, checkpoint_file, device='cuda') # test a single image and show the results img_path = './demo_test/182_7_26.png' # or img = mmcv.imread(img), which will only load it once #img=cv2.imread(img_path) #print('img_cv.shape:',img.shape) #img = mmcv.imread(img_path) #(img.shape) imgPil = Image.open(img_path) imgPil = np.array(imgPil) print('imgPil.shape:', imgPil.shape) img_pil2cv = cv2.cvtColor(imgPil, cv2.COLOR_RGB2BGR) print('img_pil2cv.shape:', img_pil2cv.shape) result = inference_segmentor(model, imgPil) # result = result[0] # result = result.astype(np.int8) # print(type(result[0][0])) # print(result) model.show_result(imgPil, result, out_file='result1126_rgb.jpg')
def main(): args = parse_args() cfg = Config.fromfile(args.config) if args.options is not None: cfg.merge_from_dict(args.options) print(f'Config:\n{cfg.pretty_text}') # dump config cfg.dump('example.py') # dump models graph if args.graph: model = init_segmentor(args.config, device='cpu') print(f'Model graph:\n{str(model)}') with open('example-graph.txt', 'w') as f: f.writelines(str(model))
def initialize(self, context): properties = context.system_properties self.map_location = 'cuda' if torch.cuda.is_available() else 'cpu' self.device = torch.device(self.map_location + ':' + str(properties.get('gpu_id')) if torch.cuda. is_available() else self.map_location) self.manifest = context.manifest model_dir = properties.get('model_dir') serialized_file = self.manifest['model']['serializedFile'] checkpoint = os.path.join(model_dir, serialized_file) self.config_file = os.path.join(model_dir, 'config.py') self.model = init_segmentor(self.config_file, checkpoint, self.device) self.initialized = True
def get_model(model_path, epoch_name): project_type = model_path.split('/')[-2] model_type = model_path.split('/')[-1] config_file = f'../../configs/{project_type}/{model_type}.py' checkpoint_file = f'/fengyouliang/model_output/mmseg_work_dirs/{project_type}/{model_type}/{epoch_name}.pth' print(f"project name: {project_type}") print(f"model type: {model_type}") print(f"loading config form {config_file}") print(f"loading checkpoint form {checkpoint_file}") print('init model right now') start = time.time() model = init_segmentor(config_file, checkpoint_file, device='cuda:0') end = time.time() print(f'loading done! time: {end - start: .4f}s') return model
def main(): args = get_parser().parse_args() img_path = args.img_path save_path = args.save_path config_file = 'configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py' checkpoint_file = 'checkpoints/pspnet_r50-d8_512x1024_40k_cityscapes_20200605_003338-2966598c.pth' model = init_segmentor(config_file, checkpoint_file, device='cuda:0') for path in tqdm(os.listdir(img_path)): if path.endswith('.jpg'): fullpath = os.path.join(img_path, path) result = inference_segmentor(model, fullpath)[0].astype(np.uint8) seg_path = os.path.join(save_path, path) cv2.imwrite(seg_path, result, [cv2.IMWRITE_JPEG_QUALITY, 75])
def main(args): url = 'http://' + args.inference_addr + '/predictions/' + args.model_name with open(args.img, 'rb') as image: tmp_res = requests.post(url, image) content = tmp_res.content if args.result_image: with open(args.result_image, 'wb') as out_image: out_image.write(content) plt.imshow(mmcv.imread(args.result_image, 'grayscale')) plt.show() else: plt.imshow(plt.imread(BytesIO(content))) plt.show() model = init_segmentor(args.config, args.checkpoint, args.device) image = mmcv.imread(args.img) result = inference_segmentor(model, image) plt.imshow(result[0]) plt.show()
def evaluate_dataset(checkpoint, device, config, data_dir="/rd-temp/mohan/iccv09Data", split_file="/rd-temp/mohan/iccv09Data/splits/val.txt", **kwargs): # initiate model model = init_segmentor(config, checkpoint, device=device) # assign config to model (Required) model.config = config result_list = [] gt_lst = [] raw_image_ids = [] with open(split_file, "r") as f: for line in f.readlines(): raw_image_ids.append(line.strip()) with tqdm(total=len(raw_image_ids)) as pbar: for image_id in raw_image_ids: # inference img = mmcv.imread(f"{data_dir}/images/{image_id}.jpg") result = inference_segmentor(model, img) # prepare image data and ground truth gt_image_path = f"{data_dir}/labels/{image_id}.regions.txt" seg_map = np.loadtxt(gt_image_path).astype(np.uint8) result_list.extend(result) gt_lst.append(seg_map) pbar.update(1) all_acc, acc, iou = mmseg.core.evaluation.mean_iou( results=result_list, gt_seg_maps=gt_lst, num_classes=config.num_classes, ignore_index=kwargs.get("ignore_index", 255), label_map=kwargs.get("label_map", {})) return (all_acc, acc, iou, iou[~np.isnan(iou)].mean())
def __init__(self, seg_net_index): self.root_split_path = TRAINING_PATH self.save_path = TRAINING_PATH + "painted_lidar/" if not os.path.exists(self.save_path): os.mkdir(self.save_path) self.seg_net_index = seg_net_index self.model = None if seg_net_index == 0: print(f'Using Segmentation Network -- {SEG_NET_OPTIONS[seg_net_index]}') self.model = torch.hub.load('pytorch/vision:v0.6.0', 'deeplabv3_resnet101', pretrained=True) self.model.eval() if torch.cuda.is_available(): self.model.to('cuda') elif seg_net_index == 1: print(f'Using Segmentation Network -- {SEG_NET_OPTIONS[seg_net_index]}') config_file = './mmseg/configs/deeplabv3plus/deeplabv3plus_r101-d8_512x1024_80k_cityscapes.py' checkpoint_file = './mmseg/checkpoints/deeplabv3plus_r101-d8_512x1024_80k_cityscapes_20200606_114143-068fcfe9.pth' self.model = init_segmentor(config_file, checkpoint_file, device='cuda:0') # TODO edit here if you want to use different device
def run_submission(model_path, epoch_name): project_type = model_path.split('/')[-2] model_type = model_path.split('/')[-1] config_file = f'../../configs/{project_type}/{model_type}.py' checkpoint_file = f'/fengyouliang/model_output/mmseg_work_dirs/{project_type}/{model_type}/{epoch_name}.pth' print(f"project name: {project_type}") print(f"model type: {model_type}") print(f"loading config form {config_file}") print(f"loading checkpoint form {checkpoint_file}") print('init model right now') model = init_segmentor(config_file, checkpoint_file, device='cuda:0') print('start inference') all_test_images = glob.glob(f"{opt.test_path}/*.tif") pbar = tqdm(all_test_images, total=len(all_test_images)) for image_file in pbar: pbar.set_description(Path(image_file).name) inference_one_image(model, image_file, project_type, model_type)
def main(): parser = ArgumentParser() parser.add_argument('img', help='Image file') parser.add_argument('config', help='Config file') parser.add_argument('checkpoint', help='Checkpoint file') parser.add_argument('--device', default='cuda:0', help='Device used for inference') parser.add_argument('--palette', default='drive', help='Color palette used for segmentation map') args = parser.parse_args() # build the model from a config file and a checkpoint file model = init_segmentor(args.config, args.checkpoint, device=args.device) # test a single image result = inference_segmentor(model, args.img) # show the results img = model.show_result(args.img, result, show=False) cv2.imwrite('out.jpg', img) show_result_pyplot(model, args.img, result)
def main(): parser = ArgumentParser() # parser.add_argument('--img', default="Image_20200925100338349.bmp", help='Image file') parser.add_argument('--img', default="star.png", help='Image file') # parser.add_argument('--img', default="demo.png", help='Image file') parser.add_argument( '--config', # default="../configs/deeplabv3/deeplabv3_r50-d8_512x1024_40k_cityscapes_custom_binary.py", default= "../configs/danet/danet_r50-d8_512x1024_40k_cityscapes_custom.py", # default="../configs/deeplabv3plus/deeplabv3plus_r101-d16-mg124_512x1024_40k_cityscapes.py", help='Config file') parser.add_argument( '--checkpoint', # default="../tools/work_dirs/deeplabv3_r50-d8_512x1024_40k_cityscapes_custom_binary/iter_200.pth", default= "../tools/work_dirs/danet_r50-d8_512x1024_40k_cityscapes_custom/iter_4000.pth", # default="../checkpoints/deeplabv3plus_r101-d16-mg124_512x1024_40k_cityscapes_20200908_005644-cf9ce186.pth", help='Checkpoint file') parser.add_argument('--device', default='cuda:0', help='Device used for inference') parser.add_argument( '--palette', default='cityscapes_custom', # default='cityscapes', help='Color palette used for segmentation map') args = parser.parse_args() # build the model from a config file and a checkpoint file model = init_segmentor(args.config, args.checkpoint, device=args.device) # test a single image result = inference_segmentor(model, args.img) # io.imsave("result.png", result[0]) # io.imshow(result[0]) # io.show() # show the results show_result_pyplot(model, args.img, result, get_palette(args.palette)) """
def test_test_time_augmentation_on_cpu(): config_file = 'configs/pspnet/pspnet_r50-d8_512x1024_40k_cityscapes.py' config = mmcv.Config.fromfile(config_file) # Remove pretrain model download for testing config.model.pretrained = None # Replace SyncBN with BN to inference on CPU norm_cfg = dict(type='BN', requires_grad=True) config.model.backbone.norm_cfg = norm_cfg config.model.decode_head.norm_cfg = norm_cfg config.model.auxiliary_head.norm_cfg = norm_cfg # Enable test time augmentation config.data.test.pipeline[1].flip = True checkpoint_file = None model = init_segmentor(config, checkpoint_file, device='cpu') img = mmcv.imread(osp.join(osp.dirname(__file__), 'data/color.jpg'), 'color') result = inference_segmentor(model, img) assert result[0].shape == (288, 512)
def inference_model(config_name, checkpoint, args, logger=None): cfg = Config.fromfile(config_name) if args.aug: if 'flip' in cfg.data.test.pipeline[ 1] and 'img_scale' in cfg.data.test.pipeline[1]: cfg.data.test.pipeline[1].img_ratios = [ 0.5, 0.75, 1.0, 1.25, 1.5, 1.75 ] cfg.data.test.pipeline[1].flip = True else: if logger is not None: logger.error(f'{config_name}: unable to start aug test') else: print(f'{config_name}: unable to start aug test', flush=True) model = init_segmentor(cfg, checkpoint, device=args.device) # test a single image result = inference_segmentor(model, args.img) # show the results if args.show: show_result_pyplot(model, args.img, result) return result
def __init__(self, args): self.args = args self.device = torch.device(args.device) # image transform input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(cfg.DATASET.MEAN, cfg.DATASET.STD), ]) # dataset and dataloader val_dataset = get_segmentation_dataset(cfg.DATASET.NAME, split='val', mode='testval', transform=input_transform) val_sampler = make_data_sampler(val_dataset, False, args.distributed) ##################### # BATCH SIZE is always 1 val_batch_sampler = make_batch_data_sampler(val_sampler, images_per_batch=cfg.TEST.BATCH_SIZE, drop_last=False) self.val_loader = data.DataLoader(dataset=val_dataset, batch_sampler=val_batch_sampler, num_workers=cfg.DATASET.WORKERS, pin_memory=True) self.classes = val_dataset.classes ### Create network ### # Segmentron model # self.model = get_segmentation_model().to(self.device) # MMSeg model mmseg_config_file = "mmseg-configs/deeplabv3plus_r101-d8_512x512_80k_ade20k.py" mmseg_pretrained = "pretrained_weights/deeplabv3plus_r101-d8_512x512_80k_ade20k_20200615_014139-d5730af7.pth" self.model = init_segmentor(mmseg_config_file, mmseg_pretrained) self.model.to(self.device) self.metric = SegmentationMetric(val_dataset.num_class, args.distributed)
color=[0, 0, 255], font_scale=1, thickness=1, back_color=[0, 0, 0]) cv2.imshow(window_name, img) global _display_inited if not _display_inited: _display_inited = True # cv2.moveWindow(window_name, -1800, 50) #home # cv2.moveWindow(window_name, -1800, 220) #office cv2.waitKey(wait) # build the model from a config file and a checkpoint file model = init_segmentor(config_file, checkpoint_file, device='cuda:0') #########tiff # for img_fn in enumerate_files("C:/_koray/test_data/space/test"): for img_fn in enumerate_files( "C:/_koray/test_data/space/SN7_buildings_test_public/test_public/L15-0369E-1244N_1479_3214_13/images_masked" ): # tiff = tiff_helper.open_tiff(img_fn) # img = cv2.imread(img_fn) # cv2.imshow("aaa", img) # cv2.waitKey(0) result = inference_segmentor(model, img_fn) # model.show_result(img_fn, result, show=True, wait_time=1000) display(False, img_fn, result, wait=1000)
height, width = 480, 640 video_name = "sample1" # height, width = 720, 1280 video_path = str(data_path / "origin" / f"{video_name}.mp4") resized_video_path = str(data_path / f"{video_name}_{height}x{width}.mp4") resized_frames_path = data_path / f"{video_name}_{height}x{width}" if not resized_frames_path.exists(): resized_frames_path.mkdir() frames_path = resized_frames_path / "img_dir" cityspaces_path = repo_path.parent / "data" / "cityscapes" device = "cuda" batch_size = 2 seg_model = init_segmentor(config_file, checkpoint_file, device=device) crf = CRF(n_spatial_dims=2, returns="log-proba").to(device) cfg = seg_model.cfg train_dataset = CityscapesDataset(data_root=cityspaces_path, pipeline=cfg.data.train.pipeline, img_dir=cfg.data.train.img_dir, ann_dir=cfg.data.train.ann_dir, test_mode=False) val_dataset = CityscapesDataset(data_root=cityspaces_path, pipeline=cfg.data.val.pipeline, img_dir=cfg.data.val.img_dir, ann_dir=cfg.data.val.ann_dir, test_mode=False)
def __init__(self, args): self.args = args self.device = torch.device(args.device) self.prefix = "ADE_cce_alpha={}".format(cfg.TRAIN.ALPHA) self.writer = SummaryWriter(log_dir=f"iccv_tensorboard/{self.prefix}") self.writer_noisy = SummaryWriter( log_dir=f"iccv_tensorboard/{self.prefix}-foggy") # image transform input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(cfg.DATASET.MEAN, cfg.DATASET.STD), ]) # dataset and dataloader train_data_kwargs = { 'transform': input_transform, 'base_size': cfg.TRAIN.BASE_SIZE, 'crop_size': cfg.TRAIN.CROP_SIZE } val_data_kwargs = { 'transform': input_transform, 'base_size': cfg.TRAIN.BASE_SIZE, 'crop_size': cfg.TEST.CROP_SIZE } train_dataset = get_segmentation_dataset(cfg.DATASET.NAME, split='train', mode='train', **train_data_kwargs) val_dataset = get_segmentation_dataset(cfg.DATASET.NAME, split='val', mode="val", **val_data_kwargs) self.classes = val_dataset.classes self.iters_per_epoch = len(train_dataset) // (args.num_gpus * cfg.TRAIN.BATCH_SIZE) self.max_iters = cfg.TRAIN.EPOCHS * self.iters_per_epoch self.ece_evaluator = ECELoss(n_classes=len(self.classes)) self.cce_evaluator = CCELoss(n_classes=len(self.classes)) train_sampler = make_data_sampler(train_dataset, shuffle=True, distributed=args.distributed) train_batch_sampler = make_batch_data_sampler(train_sampler, cfg.TRAIN.BATCH_SIZE, self.max_iters, drop_last=True) val_sampler = make_data_sampler(val_dataset, False, args.distributed) val_batch_sampler = make_batch_data_sampler(val_sampler, cfg.TEST.BATCH_SIZE, drop_last=False) self.train_loader = data.DataLoader(dataset=train_dataset, batch_sampler=train_batch_sampler, num_workers=cfg.DATASET.WORKERS, pin_memory=True) self.val_loader = data.DataLoader(dataset=val_dataset, batch_sampler=val_batch_sampler, num_workers=cfg.DATASET.WORKERS, pin_memory=True) # DEFINE data for noisy # val_dataset_noisy = get_segmentation_dataset(cfg.DATASET.NOISY_NAME, split='val', mode="val", **train_data_kwargs) # self.val_loader_noisy = data.DataLoader(dataset=val_dataset_noisy, # batch_sampler=val_batch_sampler, # num_workers=cfg.DATASET.WORKERS, # pin_memory=True) # create network # self.model = get_segmentation_model().to(self.device) mmseg_config_file = cfg.MODEL.MMSEG_CONFIG mmseg_pretrained = cfg.TRAIN.PRETRAINED_MODEL_PATH self.model = init_segmentor(mmseg_config_file, mmseg_pretrained) self.model.to(self.device) for params in self.model.backbone.parameters(): params.requires_grad = False # print params and flops if get_rank() == 0: try: show_flops_params(copy.deepcopy(self.model), args.device) except Exception as e: logging.warning('get flops and params error: {}'.format(e)) if cfg.MODEL.BN_TYPE not in ['BN']: logging.info( 'Batch norm type is {}, convert_sync_batchnorm is not effective' .format(cfg.MODEL.BN_TYPE)) elif args.distributed and cfg.TRAIN.SYNC_BATCH_NORM: self.model = nn.SyncBatchNorm.convert_sync_batchnorm(self.model) logging.info('SyncBatchNorm is effective!') else: logging.info('Not use SyncBatchNorm!') # create criterion # self.criterion = get_segmentation_loss(cfg.MODEL.MODEL_NAME, use_ohem=cfg.SOLVER.OHEM, # aux=cfg.SOLVER.AUX, aux_weight=cfg.SOLVER.AUX_WEIGHT, # ignore_index=cfg.DATASET.IGNORE_INDEX).to(self.device) self.criterion = get_segmentation_loss( cfg.MODEL.MODEL_NAME, use_ohem=cfg.SOLVER.OHEM, aux=cfg.SOLVER.AUX, aux_weight=cfg.SOLVER.AUX_WEIGHT, ignore_index=cfg.DATASET.IGNORE_INDEX, n_classes=len(train_dataset.classes), alpha=cfg.TRAIN.ALPHA).to(self.device) # optimizer, for model just includes encoder, decoder(head and auxlayer). self.optimizer = get_optimizer_mmseg(self.model) # lr scheduling self.lr_scheduler = get_scheduler(self.optimizer, max_iters=self.max_iters, iters_per_epoch=self.iters_per_epoch) # resume checkpoint if needed self.start_epoch = 0 if args.resume and os.path.isfile(args.resume): name, ext = os.path.splitext(args.resume) assert ext == '.pkl' or '.pth', 'Sorry only .pth and .pkl files supported.' logging.info('Resuming training, loading {}...'.format( args.resume)) resume_sate = torch.load(args.resume) self.model.load_state_dict(resume_sate['state_dict']) self.start_epoch = resume_sate['epoch'] logging.info('resume train from epoch: {}'.format( self.start_epoch)) if resume_sate['optimizer'] is not None and resume_sate[ 'lr_scheduler'] is not None: logging.info( 'resume optimizer and lr scheduler from resume state..') self.optimizer.load_state_dict(resume_sate['optimizer']) self.lr_scheduler.load_state_dict(resume_sate['lr_scheduler']) # evaluation metrics self.metric = SegmentationMetric(train_dataset.num_class, args.distributed) self.best_pred_miou = 0.0 self.best_pred_cces = 1e15
def main(): parser = ArgumentParser() parser.add_argument('video', help='Video file or webcam id') parser.add_argument('config', help='Config file') parser.add_argument('checkpoint', help='Checkpoint file') parser.add_argument('--device', default='cuda:0', help='Device used for inference') parser.add_argument('--palette', default='cityscapes', help='Color palette used for segmentation map') parser.add_argument('--show', action='store_true', help='Whether to show draw result') parser.add_argument('--show-wait-time', default=1, type=int, help='Wait time after imshow') parser.add_argument('--output-file', default=None, type=str, help='Output video file path') parser.add_argument('--output-fourcc', default='MJPG', type=str, help='Fourcc of the output video') parser.add_argument('--output-fps', default=-1, type=int, help='FPS of the output video') parser.add_argument('--output-height', default=-1, type=int, help='Frame height of the output video') parser.add_argument('--output-width', default=-1, type=int, help='Frame width of the output video') parser.add_argument( '--opacity', type=float, default=0.5, help='Opacity of painted segmentation map. In (0, 1] range.') args = parser.parse_args() assert args.show or args.output_file, \ 'At least one output should be enabled.' # build the model from a config file and a checkpoint file model = init_segmentor(args.config, args.checkpoint, device=args.device) # build input video cap = cv2.VideoCapture(args.video) assert (cap.isOpened()) input_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) input_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) input_fps = cap.get(cv2.CAP_PROP_FPS) # init output video writer = None output_height = None output_width = None if args.output_file is not None: fourcc = cv2.VideoWriter_fourcc(*args.output_fourcc) output_fps = args.output_fps if args.output_fps > 0 else input_fps output_height = args.output_height if args.output_height > 0 else int( input_height) output_width = args.output_width if args.output_width > 0 else int( input_width) writer = cv2.VideoWriter(args.output_file, fourcc, output_fps, (output_width, output_height), True) # start looping try: while True: flag, frame = cap.read() if not flag: break # test a single image result = inference_segmentor(model, frame) # blend raw image and prediction draw_img = model.show_result(frame, result, palette=get_palette(args.palette), show=False, opacity=args.opacity) if args.show: cv2.imshow('video_demo', draw_img) cv2.waitKey(args.show_wait_time) if writer: if draw_img.shape[0] != output_height or draw_img.shape[ 1] != output_width: draw_img = cv2.resize(draw_img, (output_width, output_height)) writer.write(draw_img) finally: if writer: writer.release() cap.release()