def main(cfg, gpu, args, progress): gpu_flag = args.gpu_flag if gpu_flag and torch.cuda.is_available(): torch.cuda.set_device(gpu) print('使用GPU进行语义分割') else: print('未开启GPU或未安装CUDA环境,设置使用CPU进行语义分割') # Network Builders if gpu_flag: 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) else: net_encoder = ModelBuilder.build_encoder( arch=cfg.MODEL.arch_encoder, fc_dim=cfg.MODEL.fc_dim, weights=cfg.MODEL.weights_encoder, gpu_flag=False) 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, gpu_flag=False) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) # Dataset and Loader dataset_test = TestDataset(cfg.list_test, cfg.DATASET) loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=cfg.TEST.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) if gpu_flag: segmentation_module.cuda() # Main loop test(segmentation_module, loader_test, gpu, gpu_flag, args, progress) print('语义分割处理完成!')
def main(args): # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(weights=args.weights_encoder) net_decoder_1 = builder.build_decoder(weights=args.weights_decoder_1) net_decoder_2 = builder.build_decoder(arch='c1',weights=args.weights_decoder_2) if args.weighted_class: crit = nn.NLLLoss(ignore_index=-1, weight=args.class_weight) else: crit = nn.NLLLoss(ignore_index=-1) # Dataset and Loader dataset_train = GTA(root=args.root_gta, cropSize=args.imgSize, is_train=1) dataset_val = CityScapes('val', root=args.root_cityscapes, cropSize=args.imgSize, max_sample=args.num_val, is_train=0) dataset_val_2 = BDD('val', root=args.root_bdd, cropSize=args.imgSize, max_sample=args.num_val, is_train=0) loader_val = torch.utils.data.DataLoader( dataset_val, batch_size=args.batch_size_eval, shuffle=False, num_workers=int(args.workers), drop_last=True) loader_val_2 = torch.utils.data.DataLoader( dataset_val_2, batch_size=args.batch_size_eval, shuffle=False, num_workers=int(args.workers), drop_last=True) args.epoch_iters = int(len(dataset_train) / args.batch_size) print('1 Epoch = {} iters'.format(args.epoch_iters)) # load nets into gpu if args.num_gpus > 1: net_encoder = nn.DataParallel(net_encoder, device_ids=range(args.num_gpus)) net_decoder_1 = nn.DataParallel(net_decoder_1, device_ids=range(args.num_gpus)) net_decoder_2 = nn.DataParallel(net_decoder_2, device_ids=range(args.num_gpus)) nets = (net_encoder, net_decoder_1, net_decoder_2, crit) for net in nets: net.cuda() history = {split: {'epoch': [], 'err': [], 'acc': [], 'mIoU': []} for split in ('train', 'val', 'val_2')} # eval evaluate(nets, loader_val, loader_val_2, history, 0, args) print('Evaluation Done!')
def main(cfg, gpu): torch.cuda.set_device(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 = TestDataset(cfg.list_test, cfg.DATASET) loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=cfg.TEST.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() # Main loop test(segmentation_module, loader_test, gpu) print('Inference done!')
def worker(args, dev_id, start_idx, end_idx, result_queue): torch.cuda.set_device(dev_id) # Dataset and Loader dataset_val = ValDataset( broden_dataset.record_list['validation'], args, max_sample=args.num_val, start_idx=start_idx, end_idx=end_idx) loader_val = torchdata.DataLoader( dataset_val, batch_size=args.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=2) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder( arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder( arch=args.arch_decoder, fc_dim=args.fc_dim, nr_classes=args.nr_classes, weights=args.weights_decoder, use_softmax=True) segmentation_module = SegmentationModule(net_encoder, net_decoder) segmentation_module.cuda() # Main loop evaluate(segmentation_module, loader_val, args, dev_id, result_queue)
def main(cfg, gpu): torch.cuda.set_device(gpu) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=cfg.MODEL.arch_encoder.lower(), fc_dim=cfg.MODEL.fc_dim, weights=cfg.MODEL.weights_encoder) net_decoder = builder.build_decoder(arch=cfg.MODEL.arch_decoder.lower(), 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_val = ValDataset(cfg.DATASET.root_dataset, cfg.DATASET.list_val, cfg.DATASET) loader_val = torchdata.DataLoader(dataset_val, batch_size=cfg.VAL.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() # Main loop evaluate(segmentation_module, loader_val, cfg, gpu) print('Evaluation Done!')
def main(args): torch.cuda.set_device(args.gpu_id) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder( arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder( arch=args.arch_decoder, fc_dim=args.fc_dim, nr_classes=args.nr_classes, weights=args.weights_decoder, use_softmax=True) segmentation_module = SegmentationModule(net_encoder, net_decoder) segmentation_module.cuda() # Dataset and Loader list_test = [{'fpath_img': args.test_img}] dataset_val = TestDataset( list_test, args, max_sample=args.num_val) loader_val = torchdata.DataLoader( dataset_val, batch_size=args.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) # Main loop test(segmentation_module, loader_val, args) print('Inference done!')
def resave(): arch_encoder = "resnet50" arch_decoder = "ppm_bilinear_deepsup" # weights_encoder="./ckpt/MacNetV2_mobilev2_scse_ppm/encoder_epoch_19.pth" weights_encoder = "" weights_decoder = "" fc_dim = 2048 builder = ModelBuilder() net_encoder = builder.build_encoder(arch=arch_encoder, fc_dim=fc_dim, weights=weights_encoder, use_softmax=True) net_decoder = builder.build_decoder(arch=arch_decoder, fc_dim=fc_dim, num_class=14, weights=weights_decoder) #net_encoder = resnet.__dict__['resnet50'](pretrained=False) net_encoder = resnet.resnet50(pretrained=False) print(net_encoder) save_path = "./ckpt/MacNetV2_mobilev2_scse_ppm/encoder_full_epoch_19.pth" torch.save(net_encoder, save_path)
def main(args): torch.cuda.set_device(args.gpu_id) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, nr_classes=args.nr_classes, weights=args.weights_decoder, use_softmax=True) segmentation_module = SegmentationModule(net_encoder, net_decoder) segmentation_module.cuda() # Dataset and Loader list_test = [{'fpath_img': args.test_img}] dataset_val = TestDataset(list_test, args, max_sample=args.num_val) loader_val = torchdata.DataLoader(dataset_val, batch_size=args.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) # Main loop test(segmentation_module, loader_val, args) print('Inference done!')
def worker(args, gpu_id, start_idx, end_idx, result_queue): torch.cuda.set_device(gpu_id) # Dataset and Loader dataset_val = ValDataset(args.list_val, args, max_sample=args.num_val, start_idx=start_idx, end_idx=end_idx) loader_val = torchdata.DataLoader(dataset_val, batch_size=args.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=2) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) segmentation_module.cuda() # Main loop evaluate(segmentation_module, loader_val, args, gpu_id, result_queue)
def main(args): torch.cuda.set_device(args.gpu_id) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, weights=args.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) # Dataset and Loader dataset_val = ValDataset(args.list_val, args, max_sample=args.num_val) loader_val = torchdata.DataLoader(dataset_val, batch_size=args.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() # Main loop evaluate(segmentation_module, loader_val, args) print('Evaluation Done!')
def worker(cfg, gpu_id, start_idx, end_idx, result_queue): torch.cuda.set_device(gpu_id) # Dataset and Loader dataset_val = ValDataset(cfg.DATASET.root_dataset, cfg.DATASET.list_val, cfg.DATASET, start_idx=start_idx, end_idx=end_idx) loader_val = torch.utils.data.DataLoader(dataset_val, batch_size=cfg.VAL.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=2) # Network Builders net_encoder = ModelBuilder.build_encoder( arch=cfg.MODEL.arch_encoder.lower(), fc_dim=cfg.MODEL.fc_dim, weights=cfg.MODEL.weights_encoder) net_decoder = ModelBuilder.build_decoder( arch=cfg.MODEL.arch_decoder.lower(), 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) segmentation_module.cuda() # Main loop evaluate(segmentation_module, loader_val, cfg, gpu_id, result_queue)
def segment(args, h, w): # absolute paths of model weights weights_encoder = os.path.join(args.model_path, 'encoder' + args.suffix) weights_decoder = os.path.join(args.model_path, 'decoder' + args.suffix) assert os.path.exists(weights_encoder) and \ os.path.exists(weights_encoder), 'checkpoint does not exitst!' builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) segmentation_module.eval() segmentation_module.cuda() # Dataset and Loader data_dict = {"in": args.content, "tar": args.style} data = load_data(data_dict, args) # Main loop with torch.no_grad(): res = test(segmentation_module, data, (h, w), args) print('Inference done!') return res
def load_model(args): args.arch_encoder = args.arch_encoder.lower() args.arch_decoder = args.arch_decoder.lower() print("Input arguments:") for key, val in vars(args).items(): print("{:16} {}".format(key, val)) # absolute paths of model weights args.weights_encoder = os.path.join(DIR_PATH, 'models', args.model_path, 'encoder' + args.suffix) args.weights_decoder = os.path.join(DIR_PATH, 'models', args.model_path, 'decoder' + args.suffix) assert os.path.exists(args.weights_encoder) and \ os.path.exists(args.weights_encoder), 'checkpoint does not exitst!' torch.cuda.set_device(args.gpu) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) segmentation_module.eval() if torch.cuda.is_available(): segmentation_module.cuda() return segmentation_module
def main(args): torch.cuda.set_device(args.gpu) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder( arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder( arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) segmentation_module.cuda() # Main loop cap = cv2.VideoCapture(0) cam_test(segmentation_module, cap, args) print('Inference done!')
def main(args): # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, segSize=args.segSize, weights=args.weights_decoder, use_softmax=True) crit = nn.NLLLoss2d(ignore_index=-1) # Dataset and Loader dataset_val = Dataset(args.list_val, args, max_sample=args.num_val, is_train=0) loader_val = torch.utils.data.DataLoader(dataset_val, batch_size=args.batch_size, shuffle=False, num_workers=2, drop_last=True) nets = (net_encoder, net_decoder, crit) for net in nets: net.cuda() # Main loop evaluate(nets, loader_val, args) print('Evaluation Done!')
def setup_test(gpu_id=0, encoder_type="resnet50dilated", decoder_type="ppm_deepsup", fc_dim=2048, model_path="baseline-resnet50dilated-ppm_deepsup", num_class=150, suffix="_epoch_20.pth"): torch.cuda.set_device(gpu_id) # absolute paths of model weights weights_encoder = os.path.join(model_path, 'encoder' + suffix) weights_decoder = os.path.join(model_path, 'decoder' + suffix) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=encoder_type, fc_dim=fc_dim, weights=weights_encoder) net_decoder = builder.build_decoder(arch=decoder_type, fc_dim=fc_dim, num_class=num_class, weights=weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) segmentation_module.cuda() segmentation_module.eval() return segmentation_module
def main(args): # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, segSize=args.segSize, num_class=args.num_class, weights=args.weights_decoder) crit = nn.NLLLoss2d(ignore_index=0) # Dataset and Loader dataset_train = ConceptDataset(args) loader_train = torch.utils.data.DataLoader( dataset_train, batch_size=args.batch_size, shuffle=True, num_workers=int(args.workers), drop_last=True) args.epoch_iters = int(len(dataset_train) / args.batch_size) print('1 Epoch = {} iters'.format(args.epoch_iters)) # load nets into gpu if args.num_gpus > 1: net_encoder = nn.DataParallel(net_encoder, device_ids=range(args.num_gpus)) net_decoder = nn.DataParallel(net_decoder, device_ids=range(args.num_gpus)) nets = (net_encoder, net_decoder, crit) if args.num_gpus > 0: for net in nets: net.cuda() # Set up optimizers optimizers = create_optimizers(nets, args) # Main loop history = {split: {'epoch': [], 'err': [], 'acc': []} for split in ('train', 'val')} # initial eval # evaluate(nets, loader_val, history, 0, args) for epoch in range(1, args.num_epoch + 1): train(nets, loader_train, optimizers, history, epoch, args) # Evaluation and visualization # if epoch % args.eval_epoch == 0: # evaluate(nets, loader_val, history, epoch, args) # checkpointing checkpoint(nets, history, args) # adjust learning rate adjust_learning_rate(optimizers, epoch, args) print('Training Done!')
def fetch_teacher_outputs(args): teacher_outputs = [] # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) # dataset and loader, use train dataset_train = TrainDataset(args.list_train, args, batch_per_gpu=args.batch_size_per_gpu) loader_train = torchdata.DataLoader( dataset_train, batch_size=args.num_gpus, # we have modified data_parallel shuffle=False, # we do not use this param collate_fn=user_scattered_collate, num_workers=int(args.workers), drop_last=True, pin_memory=True) segmentation_module.cuda() # here may be some problems for i, batch_data in enumerate(loader_train): batch_data = batch_data[0] # get data list 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]) # 预测输出的形状 pred = torch.zeros(1, args.num_class, segSize[0], segSize[1]) 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_id) # forward pass pred_tmp = segmentation_module(feed_dict, segSize=segSize) pred = pred + pred_tmp.cpu() / len(args.imgSize) _, preds = torch.max(pred, dim=1) preds.as_numpy(preds.squeeze(0)) teacher_outputs.append(preds) return teacher_outputs
def main(args): # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) nr_classes = broden_dataset.nr.copy() nr_classes['part'] = sum( [len(parts) for obj, parts in broden_dataset.object_part.items()]) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, nr_classes=nr_classes, weights=args.weights_decoder) # TODO(LYC):: move criterion outside model. # crit = nn.NLLLoss(ignore_index=-1) if args.arch_decoder.endswith('deepsup'): segmentation_module = SegmentationModule(net_encoder, net_decoder, args.deep_sup_scale) else: segmentation_module = SegmentationModule(net_encoder, net_decoder) print('1 Epoch = {} iters'.format(args.epoch_iters)) # create loader iterator iterator_train = create_multi_source_train_data_loader(args=args) # load nets into gpu if args.num_gpus > 1: segmentation_module = UserScatteredDataParallel(segmentation_module, device_ids=range( args.num_gpus)) # For sync bn patch_replication_callback(segmentation_module) segmentation_module.cuda() # Set up optimizers nets = (net_encoder, net_decoder) optimizers = create_optimizers(nets, args) # Main loop history = {'train': {'epoch': [], 'loss': [], 'acc': []}} for epoch in range(args.start_epoch, args.num_epoch + 1): train(segmentation_module, iterator_train, optimizers, history, epoch, args) # checkpointing checkpoint(nets, history, args, epoch) print('Training Done!')
def main(args): # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder_1 = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder_1) net_decoder_2 = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder_2) net_syn = builder.build_syn(weights=args.weights_syn) crit = nn.NLLLoss2d(ignore_index=-1) # Dataset and Loader # dataset_val = CityScapes('val', root=args.root_cityscapes, max_sample=args.num_val, is_train=0) dataset_val = BDD('val', root=args.root_unlabeled, cropSize=args.imgSize, max_sample=args.num_val, is_train=0) loader_val = torch.utils.data.DataLoader(dataset_val, batch_size=args.batch_size, shuffle=False, num_workers=2, drop_last=True) nets = (net_encoder, net_decoder_1, net_decoder_2, net_syn, crit) for net in nets: net.cuda() # Main loop evaluate(nets, loader_val, args) print('Evaluation Done!')
def main(args): # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder( arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder( arch=args.arch_decoder, fc_dim=args.fc_dim, nr_classes=args.nr_classes, weights=args.weights_decoder) # TODO(LYC):: move criterion outside model. # crit = nn.NLLLoss(ignore_index=-1) if args.arch_decoder.endswith('deepsup'): segmentation_module = SegmentationModule( net_encoder, net_decoder, args.deep_sup_scale) else: segmentation_module = SegmentationModule( net_encoder, net_decoder) print('1 Epoch = {} iters'.format(args.epoch_iters)) # create loader iterator iterator_train = create_multi_source_train_data_loader(args=args) # load nets into gpu if args.num_gpus > 1: segmentation_module = UserScatteredDataParallel( segmentation_module, device_ids=range(args.num_gpus)) # For sync bn patch_replication_callback(segmentation_module) segmentation_module.cuda() # Set up optimizers nets = (net_encoder, net_decoder) optimizers = create_optimizers(nets, args) # Main loop history = {'train': {'epoch': [], 'loss': [], 'acc': []}} for epoch in range(args.start_epoch, args.num_epoch + 1): train(segmentation_module, iterator_train, optimizers, history, epoch, args) # checkpointing checkpoint(nets, history, args, epoch) print('Training Done!')
def main(args): torch.cuda.set_device(args.gpu) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder( arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder( arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) # Dataset and Loader # list_test = [{'fpath_img': args.test_img}] # test_chk = [] # testing = os.listdir("/home/teai/externalhd2/BDD100K/segmentation_v2/test/") # for i in testing: # if(i.endswith(".jpg")): # test_chk.append("/home/teai/externalhd2/BDD100K/segmentation_v2/test/"+i) test_chk = ['./05b07068-373666cb.jpg'] print(type(args.test_imgs)) list_test = [{'fpath_img': x} for x in test_chk] #list_test=[{'fpath_img': 'frame_143.png'},{'fpath_img': 'frame_100.png'},{'fpath_img': 'frame_1.png'}] #print("list_test",list_test) dataset_test = TestDataset( list_test, args, max_sample=args.num_val) loader_test = torchdata.DataLoader( dataset_test, batch_size=args.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() # Main loop start=time.time() test(segmentation_module, loader_test, args) end=time.time()
def main(self): torch.cuda.set_device(self.gpu) # Network Builders net_encoder = ModelBuilder.build_encoder( arch=self.cfg.MODEL.arch_encoder, fc_dim=self.cfg.MODEL.fc_dim, weights=self.cfg.MODEL.weights_encoder) net_decoder = ModelBuilder.build_decoder( arch=self.cfg.MODEL.arch_decoder, fc_dim=self.cfg.MODEL.fc_dim, num_class=self.cfg.DATASET.num_class, weights=self.cfg.MODEL.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) self.segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) self.segmentation_module.cuda() self.seg_pub = rospy.Publisher(self.img_out, sensor_msgs.msg.Image, queue_size=1) rospy.Subscriber(self.img_in, sensor_msgs.msg.Image, self.image_callback) rospy.loginfo("Listening for image messages on topic %s..." % self.img_in) rospy.loginfo("Publishing segmented images to topic %s..." % self.img_out) rospy.loginfo("Waiting for loader from queue...") while not rospy.is_shutdown(): rospy.sleep(0.01) try: loader = self.loader_q.get_nowait() self.ready = False self.run_inference(loader) self.ready = True except queue.Empty: pass rospy.spin()
def main(args): # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, segSize=args.segSize, weights=args.weights_decoder, use_softmax=True) nets = (net_encoder, net_decoder) for net in nets: net.cuda() # single pass test(nets, args)
def inference(cfg, image_path, gpu=0): torch.cuda.set_device(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 = TestDataset( cfg.list_test, cfg.DATASET) ''' dataset_test = TestROPRidgeDataset( root_dataset=cfg.DATASET.root_dataset, opt=cfg.DATASET, img_folder=cfg.DATASET.img_folder_val, annotation_folder="annotations", anno_filename=cfg.DATASET.list_val, batch_per_gpu=cfg.TRAIN.batch_size_per_gpu) loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=cfg.TEST.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() # Main loop test(segmentation_module, image_path, gpu) print('Inference done!')
def main(args): torch.cuda.set_device(args.gpu_id) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) # Dataset and Loader # list_test = [{'fpath_img': args.test_img}] # Pass a single image if args.test_imgs is not None: list_test = [{'fpath_img': x} for x in args.test_imgs] if args.test_img_path is not None: list_test = [{ 'fpath_img': os.path.join(args.test_img_path, x) } for x in os.listdir(args.test_img_path)] list_test = list_test[:20] dataset_test = TestDataset(list_test, args, max_sample=args.num_val) loader_test = torchdata.DataLoader(dataset_test, batch_size=args.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() # Main loop test(segmentation_module, loader_test, args) print('Inference done!')
def _prepare_module(self): net_encoder = ModelBuilder.build_encoder( arch=self.cfg.MODEL.arch_encoder, fc_dim=self.cfg.MODEL.fc_dim, weights=self.cfg.MODEL.weights_encoder) net_decoder = ModelBuilder.build_decoder( arch=self.cfg.MODEL.arch_decoder, fc_dim=self.cfg.MODEL.fc_dim, num_class=self.cfg.DATASET.num_class, weights=self.cfg.MODEL.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) segmentation_module.cuda() return segmentation_module
def module_init(cfg): # 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) # NLLLoss crit = nn.NLLLoss(ignore_index=-1) # Instantiate segmentation module segmentation_module = SegmentationModule(net_encoder, net_decoder, crit).cuda() # Evaluation mode segmentation_module.eval() return segmentation_module
def loadModel(model): net_encoder = ModelBuilder.build_encoder(arch=model["encoder_arch"], fc_dim=model["fc_dim"], weights=os.path.join( cwd, model["encoder_weights"])) net_decoder = ModelBuilder.build_decoder(arch=model["decoder_arch"], fc_dim=model["fc_dim"], num_class=model["num_class"], weights=os.path.join( cwd, model["decoder_weights"]), use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) segmentation_module.cuda() return segmentation_module
def main(args): # download data from buckets #download_from_buckets(args.test_img) bucket_name = "ubiquity-kube-mlengine-pytorch_trial" source_blob_name = "ade20k/davidgarrett.jpg" destination_file_name = "/davidgarrett.jpg" download_blob(bucket_name, source_blob_name, destination_file_name) args.test_img = destination_file_name bucket_url = args.weights_encoder print(bucket_url, type(bucket_url)) destination_folder = 'model_path_en' args.weights_encoder = download_from_buckets(bucket_url, destination_folder) bucket_url = args.weights_decoder print(bucket_url, type(bucket_url)) destination_folder = 'model_path_de' args.weights_decoder = download_from_buckets(bucket_url, destination_folder) print('gsutil works') # Network Builders writer = SummaryWriter() builder = ModelBuilder() net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, segSize=args.segSize, weights=args.weights_decoder, use_softmax=True) nets = (net_encoder, net_decoder) if torch.cuda.is_available(): for net in nets: net.cuda() # single pass test(nets, args) print('Done! Output is saved in {}'.format(args.result))
def main(args): torch.cuda.set_device(args.gpu) # Network Builders builder = ModelBuilder() net_encoder = builder.build_encoder( arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder( arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=args.num_class, weights=args.weights_decoder, use_softmax=True) crit = nn.NLLLoss(ignore_index=-1) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) # Dataset and Loader if len(args.test_imgs) == 1 and os.path.isdir(args.test_imgs[0]): test_imgs = find_recursive(args.test_imgs[0]) else: test_imgs = args.test_imgs list_test = [{'fpath_img': x} for x in test_imgs] dataset_test = TestDataset( list_test, args, max_sample=args.num_val) loader_test = torchdata.DataLoader( dataset_test, batch_size=args.batch_size, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) segmentation_module.cuda() # Main loop test(segmentation_module, loader_test, args) print('Inference done!')
def main(args): torch.cuda.set_device(args.gpu) crit = nn.NLLLoss(ignore_index=-1) # Dataset and Loader dataset_val = Dataset(args, split_name=args.split_name, batch_per_gpu=args.batchsize) loader_val = torchdata.DataLoader(dataset_val, batch_size=args.batchsize, shuffle=False, collate_fn=user_scattered_collate, num_workers=5, drop_last=True) # Network Builders builder = ModelBuilder() print('Loading encoder from: %s' % (args.weights_encoder)) print('Loading decoder from: %s' % (args.weights_decoder)) net_encoder = builder.build_encoder(arch=args.arch_encoder, fc_dim=args.fc_dim, weights=args.weights_encoder) net_decoder = builder.build_decoder(arch=args.arch_decoder, fc_dim=args.fc_dim, num_class=dataset_val.num_classes, weights=args.weights_decoder, use_softmax=True) segmentation_module = SegmentationModule(net_encoder, net_decoder, crit) segmentation_module.cuda() # Main loop if args.dataset.startswith('nyuv2sn'): metrics = evaluate_surface_normals(segmentation_module, loader_val, args) else: metrics = evaluate_segmentation(segmentation_module, loader_val, args) save_pickle(args.result_file, metrics) print('Evaluation Done!')