def _load_mistcsail_model(encoder, decoder, encoder_path=None, decoder_path=None): from mit_semseg.models import ModelBuilder, SegmentationModule from mit_semseg.config import cfg if encoder_path is None: encoder_path = str(Path(__file__).parents[0].joinpath('weights', f'encoder_{encoder}.pth')) if decoder_path is None: decoder_path = str(Path(__file__).parents[0].joinpath('weights', f'decoder_{decoder}.pth')) config_file = str(Path(__file__).parents[0].joinpath('config', f'ade20k-{encoder}-{decoder}.yaml')) cfg.merge_from_file(config_file) net_encoder = ModelBuilder.build_encoder( arch=encoder, fc_dim=cfg.MODEL.fc_dim, weights=encoder_path) net_decoder = ModelBuilder.build_decoder( arch=decoder, fc_dim=cfg.MODEL.fc_dim, num_class=cfg.DATASET.num_class, weights=decoder_path, use_softmax=True) crit = torch.nn.NLLLoss(ignore_index=-1) model = SegmentationModule(net_encoder, net_decoder, crit) return model
def predict(imgPath): args = Namespace(imgs=imgPath, cfg="config/ade20k-resnet50dilated-ppm_deepsup.yaml", opts=None) 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) # assert os.path.exists(cfg.MODEL.weights_encoder) and \ # os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!" imgs = [args.imgs] assert len(imgs), "imgs should be a path to image (.jpg) or directory." cfg.list_test = [{'fpath_img': x} for x in imgs] if not os.path.isdir(cfg.TEST.result): os.makedirs(cfg.TEST.result) prediction = main(cfg, None) return prediction
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 _load_model(self, cfg_file, gpu): if gpu is not None: torch.cuda.set_device(0) basepath = rospkg.RosPack().get_path('object_segmentation') cfg.merge_from_file(basepath + "/" + cfg_file) logger = setup_logger(distributed_rank=0) logger.info(f"Loaded configuration file {cfg_file}") logger.info("Running with config:\n{}".format(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 = ( Path(basepath) / cfg.DIR / ('encoder_' + cfg.TEST.checkpoint)).as_posix() cfg.MODEL.weights_decoder = ( Path(basepath) / cfg.DIR / ('decoder_' + cfg.TEST.checkpoint)).as_posix() if not os.path.exists(cfg.MODEL.weights_encoder) or not os.path.exists( cfg.MODEL.weights_decoder): download.ycb(Path(basepath) / 'ckpt') assert os.path.exists( cfg.MODEL.weights_encoder ), f"checkpoint {cfg.MODEL.weights_encoder} does not exitst!" assert os.path.exists( cfg.MODEL.weights_decoder ), f"checkpoint {cfg.MODEL.weights_decoder} does not exitst!" # 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) if self.gpu is not None: segmentation_module.cuda() segmentation_module.eval() self.model = segmentation_module
def main(): args = build_argparser().parse_args() cfg.merge_from_file(args.cfg) #cfg.merge_from_list(args.opts) logger = setup_logger(distributed_rank=0) # TODO logger.info("Loaded configuration file {}".format(args.cfg)) logger.info("Running with config:\n{}".format(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) assert os.path.exists(cfg.MODEL.weights_encoder) and \ os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!" # 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) seg_model = SegmentationModule(net_encoder, net_decoder, None) input = torch.randn(1, 3, 320, 320) #input_decoder = torch.randn(1, 180, 720, 3, 3) segSize = (256, 256) #torch.onnx.export(net_encoder, input, "encoder.onnx", opset_version=9, verbose=True) torch.onnx.export(seg_model, input, "seg.onnx", opset_version=9, verbose=True)
def main(): args = build_argparser().parse_args() cfg.merge_from_file(args.cfg) #cfg.merge_from_list(args.opts) logger = setup_logger(distributed_rank=0) # TODO logger.info("Loaded configuration file {}".format(args.cfg)) logger.info("Running with config:\n{}".format(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) assert os.path.exists(cfg.MODEL.weights_encoder) and \ os.path.exists(cfg.MODEL.weights_decoder), "checkpoint does not exitst!" # 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) ''' input = torch.randn(1, 3, 320, 320) input_decoder = torch.randn(1, 512, 2048, 40, 40) print(input_decoder[-1]) segSize = (256, 256)
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
help="path to config file", type=str, ) parser.add_argument("--gpu", default=0, type=int, help="gpu id for evaluation") parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg.merge_from_file(args.cfg) cfg.merge_from_list(args.opts) # cfg.freeze() logger = setup_logger(distributed_rank=0) # TODO logger.info("Loaded configuration file {}".format(args.cfg)) logger.info("Running with config:\n{}".format(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)
parser.add_argument("--input", default=f'{TEST_DIR}\\input.mp4', help="Path to input video file") parser.add_argument("--output", default=f'{TEST_DIR}\\output.mp4', help="Path to output style video file") parser.add_argument("--label", default="car", help="Path to output style video file") parser.add_argument("--use-cpu", action='store_true', help="Use CPU instead of GPU") parser.add_argument("--gpu-device", type=int, default=None, help="GPU device index") parser.add_argument("--fps", type=int, default=None, help="FPS of output video") args = parser.parse_args() label_index = -1 with open(LABEL_CSV, 'r') as csv_file: csv_reader = csv.reader(csv_file) for row in csv_reader: if label_index > 0: break for column in row: if (args.label in column): label_index = int(row[0]) - 1 style_model = StyleModel(DEFAULT_MODEL, use_gpu=not args.use_cpu, gpu_device=args.gpu_device) cfg.merge_from_file("config/ade20k-resnet50dilated-ppm_deepsup.yaml") segmentation_model = SegmentationModel( DEFAULT_MODEL, use_gpu=not args.use_cpu, gpu_device=args.gpu_device, cfg=cfg) fix_path = lambda x: x.replace('\\', r'\\') reader = VideoReader(fix_path(args.input), fps=args.fps) output_writer = VideoWriter(fix_path(args.output), reader.width, reader.height, reader.fps) with output_writer: frame_number = 0 for frame in reader: print(f"Processing frame {frame_number}") frame_number+=1 image = np.array(frame) style = style_model.run(image)
def inference_prob( img, device, select_model_option="ade20k-resnet50dilated-ppm_deepsup" ): # select_model_option = "ade20k-mobilenetv2dilated-c1_deepsup" / "ade20k-hrnetv2" '''Load the data and preprocess settings Input: img - the path of our target image device - Current device running select_model_option - name of NN we use ''' cfg_ss.merge_from_file("ss/config/" + select_model_option + ".yaml") logger = setup_logger(distributed_rank=0) # TODO cfg_ss.MODEL.arch_encoder = cfg_ss.MODEL.arch_encoder.lower() cfg_ss.MODEL.arch_decoder = cfg_ss.MODEL.arch_decoder.lower() # absolute paths of model weights cfg_ss.MODEL.weights_encoder = os.path.join( 'ss/' + cfg_ss.DIR, 'encoder_' + cfg_ss.TEST.checkpoint) cfg_ss.MODEL.weights_decoder = os.path.join( 'ss/' + cfg_ss.DIR, 'decoder_' + cfg_ss.TEST.checkpoint) assert os.path.exists(cfg_ss.MODEL.weights_encoder) and os.path.exists( cfg_ss.MODEL.weights_decoder), "checkpoint does not exist!" # generate testing image list imgs = [img] assert len(imgs), "imgs should be a path to image (.jpg) or directory." cfg_ss.list_test = [{'fpath_img': x} for x in imgs] if not os.path.isdir(cfg_ss.TEST.result): os.makedirs(cfg_ss.TEST.result) if torch.cuda.is_available(): torch.cuda.set_device(device) # Network Builders net_encoder = ModelBuilder.build_encoder( arch=cfg_ss.MODEL.arch_encoder, fc_dim=cfg_ss.MODEL.fc_dim, weights=cfg_ss.MODEL.weights_encoder) net_decoder = ModelBuilder.build_decoder( arch=cfg_ss.MODEL.arch_decoder, fc_dim=cfg_ss.MODEL.fc_dim, num_class=cfg_ss.DATASET.num_class, weights=cfg_ss.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 = TestDataset(cfg_ss.list_test, cfg_ss.DATASET) loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=cfg_ss.TEST.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.to(device) # Main loop return segmentation_module, loader_test