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")
