def __init__(self, args): self.args = args # output folder self.outdir = 'test_result' if not os.path.exists(self.outdir): os.makedirs(self.outdir) # image transform input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([.485, .456, .406], [.229, .224, .225]), ]) # dataset and dataloader val_dataset = get_segmentation_dataset(args.dataset, split='val', mode='test', transform=input_transform) self.val_loader = data.DataLoader(dataset=val_dataset, batch_size=1, shuffle=False) # create network self.model = get_fast_scnn(args.dataset, aux=args.aux, pretrained=True, root=args.save_folder).to(args.device) print('Finished loading model!') self.metric = SegmentationMetric(val_dataset.num_class)
def demo(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # output folder if not os.path.exists(args.outdir): os.makedirs(args.outdir) # image transform transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) image = Image.open(args.input_pic).convert('RGB') image = transform(image).unsqueeze(0).to(device) model = get_fast_scnn(args.dataset, pretrained=True, root=args.weights_folder).to(device) print('Finished loading model!') model.eval() with torch.no_grad(): outputs = model(image) pred = torch.argmax(outputs[0], 1).squeeze(0).cpu().data.numpy() mask = get_color_pallete(pred, args.dataset) outname = os.path.splitext(os.path.split(args.input_pic)[-1])[0] + '.png' mask.save(os.path.join(args.outdir, outname))
def demo(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # output folder if not os.path.exists(args.outdir): os.makedirs(args.outdir) # image transform transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) # image = Image.open(args.input_pic).convert('RGB') image = Image.open(args.input_pic).resize((2048, 1024)).convert('RGB') # img_name = os.path.splitext(os.path.split(args.input_pic)[-1])[0] + '.jpg' # image.save(os.path.join(args.outdir, img_name)) print('image.shape:', image.size) # (2048, 1024) image = transform(image).unsqueeze(0).to(device) print('image.shape2:', image.shape) # [1, 3, 1024, 2048] model = get_fast_scnn(args.dataset, pretrained=True, root=args.weights_folder, map_cpu=args.cpu).to(device) # pretrained=True print('Finished loading model!') model.eval() with torch.no_grad(): start = time.time() outputs = model(image) print('outputs:', outputs[0].shape) # [1, 19, 1024, 2048] print('cost_time:', str(time.time() - start)) pred = torch.argmax(outputs[0], 1).squeeze(0).cpu().data.numpy() print('pred.shape:', pred.shape) print('pred:', pred) mask = get_color_pallete(pred, args.dataset) outname = os.path.splitext(os.path.split(args.input_pic)[-1])[0] + '.png' mask.save(os.path.join(args.outdir, outname))
def __init__(self): rospy.loginfo("Changing all ROS topics to cv image") #self.image_sub=rospy.Subscriber("/image_publisher_1559912716887579484/image_raw",Image,self.imgcallback) #self.image_sub=rospy.Subscriber("/cv_camera/image_raw",Image,self.imgcallback) self.image_sub = rospy.Subscriber("/zed/left/image_raw_color", Image, self.imgcallback) self.bridge = CvBridge() self.image_data = None self.cv_image = None rospy.loginfo("All objects for ros to cv conversion initialised") self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.model = get_fast_scnn(args.dataset, pretrained=True, root=args.weights_folder, map_cpu=args.cpu).to(self.device) print('Finished loading model!') self.model.eval() self.transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) self.mtx = np.array([[709.103066, 0.000000, 621.543559], [0.000000, 709.978057, 333.677376], [0.000000, 0.000000, 1.000000]]) print(self.mtx.shape) self.dist = np.array( [-0.163186, 0.026619, 0.000410, 0.000569, 0.000000]) rospy.sleep(4)
def __init__(self,NUM_CHANNELS): MODEL_NAME = "fast_scnn" if NUM_CHANNELS==4 else "fast_scnn_rgb" self.model_name = MODEL_NAME self.outdir = ARTIFACT_DETECTION_DIR + "/inference/" + self.model_name if not os.path.exists(self.outdir): os.makedirs(self.outdir) if MODEL_NAME == "fast_scnn": input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([.485, .456, .406,0.4], [.229, .224, .225,0.4]), ]) else: input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([.485, .456, .406], [.229, .224, .225]), ]) val_dataset = get_segmentation_dataset(num_channels=NUM_CHANNELS,root=ARGS_INFERENCE_DIR, split='val', mode='testval', transform=input_transform) self.val_loader = data.DataLoader(dataset=val_dataset, batch_size=1, shuffle=False) WEIGHTS_PATH = ARTIFACT_DETECTION_DIR+"/weights/fast_scnn_rgb" if NUM_CHANNELS==3 else ARTIFACT_DETECTION_DIR+"/weights/fast_scnn" print("Weights Path:", WEIGHTS_PATH) self.model = get_fast_scnn("fast_scnn",num_channels=NUM_CHANNELS, aux=False, pretrained=True, root=WEIGHTS_PATH).cuda() self.metric = SegmentationMetric(val_dataset.num_class)
def __init__(self, args): self.args = args # image transform input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([.485, .456, .406], [.229, .224, .225]), ]) # dataset and dataloader data_kwargs = {'transform': input_transform, 'base_size': args.base_size, 'crop_size': args.crop_size} train_dataset = get_segmentation_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs) val_dataset = get_segmentation_dataset(args.dataset, split='val', mode='val', **data_kwargs) self.train_loader = data.DataLoader(dataset=train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True) self.val_loader = data.DataLoader(dataset=val_dataset, batch_size=1, shuffle=False) # create network self.model = get_fast_scnn(dataset=args.dataset, aux=args.aux) if torch.cuda.device_count() > 1: self.model = torch.nn.DataParallel(self.model, device_ids=[0, 1, 2]) self.model.to(args.device) # resume checkpoint if needed if args.resume: if os.path.isfile(args.resume): name, ext = os.path.splitext(args.resume) assert ext == '.pkl' or '.pth', 'Sorry only .pth and .pkl files supported.' print('Resuming training, loading {}...'.format(args.resume)) self.model.load_state_dict(torch.load(args.resume, map_location=lambda storage, loc: storage)) # create criterion self.criterion = MixSoftmaxCrossEntropyOHEMLoss(aux=args.aux, aux_weight=args.aux_weight, ignore_index=-1).to(args.device) # optimizer self.optimizer = torch.optim.SGD(self.model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # lr scheduling self.lr_scheduler = LRScheduler(mode='poly', base_lr=args.lr, nepochs=args.epochs, iters_per_epoch=len(self.train_loader), power=0.9) # evaluation metrics self.metric = SegmentationMetric(train_dataset.num_class) self.best_pred = 0.0
def __init__(self, args, cfg=None): # train_dataset = [build_dataset(cfg.data.train)] # self.dataset= train_dataset # val_dataset = [build_dataset(cfg.data.test)] # if len(cfg.workflow) == 2: # train_dataset.append(build_dataset(cfg.data.val)) # train_data_loaders = [ # build_dataloader( # ds, # cfg.data.imgs_per_gpu, # cfg.data.workers_per_gpu, # # cfg.gpus, # dist=False) for ds in train_dataset # ] # val_data_loader = build_dataloader( # val_dataset, # imgs_per_gpu=1, # workers_per_gpu=cfg.data.workers_per_gpu, # dist=False, # shuffle=False) # self.train_loader = train_data_loaders[0] # self.val_loader = val_data_loader self.args = args # image transform input_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]), ]) # dataset and dataloader data_kwargs = { 'transform': input_transform, 'base_size': args.base_size, 'crop_size': args.crop_size } train_dataset = get_segmentation_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs) val_dataset = get_segmentation_dataset(args.dataset, split='val', mode='val', **data_kwargs) self.train_loader = data.DataLoader(dataset=train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True) self.val_loader = data.DataLoader(dataset=val_dataset, batch_size=1, shuffle=False) # create network self.model = get_fast_scnn(dataset=args.dataset, aux=args.aux) if torch.cuda.device_count() > 1: self.model = torch.nn.DataParallel(self.model, device_ids=[0, 1, 2]) self.model.to(args.device) # resume checkpoint if needed if args.resume: if os.path.isfile(args.resume): name, ext = os.path.splitext(args.resume) assert ext == '.pkl' or '.pth', 'Sorry only .pth and .pkl files supported.' print('Resuming training, loading {}...'.format(args.resume)) self.model.load_state_dict( torch.load(args.resume, map_location=lambda storage, loc: storage)) # create criterion self.criterion = MixSoftmaxCrossEntropyOHEMLoss( aux=args.aux, aux_weight=args.aux_weight, ignore_index=-1).to(args.device) # optimizer self.optimizer = torch.optim.SGD(self.model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # lr scheduling self.lr_scheduler = LRScheduler(mode='poly', base_lr=args.lr, nepochs=args.epochs, iters_per_epoch=len(self.train_loader), power=0.9) # evaluation metrics self.metric = SegmentationMetric(train_dataset.num_class) self.best_pred = 0.0
# 海陆分割模型初始化 dataset = 'citys' weights_folder = 'yasuoweight' cpu = False device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 海陆分割输入图片需要的一些处理,转tensor和归一化 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) # 海陆分割模型加载 start4 = time.time() land_sea_model = get_fast_scnn(dataset, pretrained=True, root=weights_folder, map_cpu=cpu).to(device) land_sea_model.eval() end4 = time.time() print('time_load:', end4 - start4) # pycuda之后要用到的一个逐元素操作函数 gpu_thresh = ElementwiseKernel("float *in, float *out", "out[i] = (in[i]>=10)*255;", "gpu_thresh") gpu_absthresh = ElementwiseKernel("float *in, float *out", "out[i] = (((in[i]>0)*in[i])>=10)*255;", "gpu_absthresh") gpu_sub = ElementwiseKernel("float *in1, float *in2,float *out", "out[i] = in1[i]-in2[i];", "gpu_sub")