def main():
global args, best_EPE, save_path
args = parser.parse_args()
save_path = '{},{},{}epochs{},b{},lr{}'.format(
args.arch, args.solver, args.epochs,
',epochSize' + str(args.epoch_size) if args.epoch_size > 0 else '',
args.batch_size, args.lr)
if not args.no_date:
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
save_path = os.path.join(timestamp, save_path)
save_path = os.path.join(args.dataset, save_path)
print('=> will save everything to {}'.format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
train_writer = SummaryWriter(os.path.join(save_path, 'train'))
test_writer = SummaryWriter(os.path.join(save_path, 'test'))
output_writers = []
for i in range(3):
output_writers.append(
SummaryWriter(os.path.join(save_path, 'test', str(i))))
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0], std=[args.div_flow, args.div_flow])
])
print("=> fetching img pairs in '{}'".format(args.data))
test_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
target_transform=target_transform)
print('{} samples found'.format(len(test_set)))
val_loader = torch.utils.data.DataLoader(test_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=False)
# create model
if args.pretrained:
network_data = torch.load(args.pretrained)
args.arch = network_data['arch']
print("=> using pre-trained model '{}'".format(args.arch))
else:
network_data = None
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch](network_data).cuda()
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
best_EPE = validate(val_loader, model, 0, output_writers)
return
def __getitem__(self, index):
img_idx = self.lists[index]
img = self.imgs[img_idx].transpose(2, 1, 0)
#img = self.imgs[img_idx]
dpt = self.dpts[img_idx].transpose(1, 0)
#dpt = self.dpts[img_idx]
#image = Image.fromarray(np.uint8(img))
#depth = Image.fromarray(np.uint8(dpt))
#image.save('img1.png')
#input_transform = transforms.Compose([flow_transforms.Scale(228)])
#input_transform = transforms.Compose([flow_transforms.ArrayToTensor()])
input_transform = transforms.Compose(
[flow_transforms.Scale(228),
flow_transforms.ArrayToTensor()])
#target_depth_transform = transforms.Compose([flow_transforms.Scale(228)])
#target_depth_transform = transforms.Compose([flow_transforms.ArrayToTensor()])
target_depth_transform = transforms.Compose([
flow_transforms.Scale_Single(228),
flow_transforms.ArrayToTensor()
])
img = input_transform(img)
dpt = target_depth_transform(dpt)
#image = Image.fromarray(np.uint8(img))
#image.save('img2.png')
return img, dpt
def __init__(self,
args,
root='datasets/FlyingChairs_release/data',
mode='train'):
# Normalize images to [-1,1]
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
target_transform = transforms.Compose(
[flow_transforms.ArrayToTensor()])
# Simple Aug
if mode == 'train':
co_transform = flow_transforms.Compose([
flow_transforms.RandomTranslate(cfg.RANDOM_TRANS),
flow_transforms.RandomCrop(
(cfg.CROP_SIZE[0], cfg.CROP_SIZE[1])),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
else:
co_transform = None
self.root = root
self.transform = input_transform
self.target_transform = target_transform
self.co_transform = co_transform
images = []
for flow_map in sorted(glob.glob(os.path.join(root, '*_flow.flo'))):
flow_map = os.path.basename(flow_map)
root_filename = flow_map[:-9]
img1 = root_filename + '_img1.ppm'
img2 = root_filename + '_img2.ppm'
if not (os.path.isfile(os.path.join(dir, img1))
and os.path.isfile(os.path.join(dir, img2))):
continue
images.append([[img1, img2], flow_map])
train_list, test_list = split2list(root, 'FlyingChairs_train_val.txt')
if mode == 'train':
self.path_list = train_list
else:
self.path_list = test_list
def main():
global args, save_path
args = parser.parse_args()
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
if args.output is None:
save_path = data_dir/'flow'
else:
save_path = Path(args.output)
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255]),
transforms.Normalize(mean=[0.411,0.432,0.45], std=[1,1,1])
])
img_pairs = []
for ext in args.img_exts:
test_files = data_dir.files('*0.{}'.format(ext))
for file in test_files:
img_pair = file.parent / (file.namebase[:-1] + '1.{}'.format(ext))
if img_pair.isfile():
img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = input_transform(imread(img1_file))
img2 = input_transform(imread(img2_file))
input_var = torch.cat([img1, img2]).unsqueeze(0)
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var)
if args.upsampling is not None:
output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False)
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
rgb_flow = flow2rgb(args.div_flow * flow_output, max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0)
imsave(save_path/'{}{}.png'.format(img1_file.namebase[:-1], suffix), to_save)
def main():
global args
args = parser.parse_args()
test_list = make_dataset(args.data)
flow_epe = AverageMeter()
avg_parts_epe = {}
for bk in BODY_MAP.keys():
avg_parts_epe[bk] = AverageMeter()
for i, (pred_path, flow_path, seg_path) in enumerate(tqdm(test_list)):
predflow = flow_transforms.ArrayToTensor()(load_flo(pred_path))
gtflow = flow_transforms.ArrayToTensor()(load_flo(flow_path))
segmask = flow_transforms.ArrayToTensor()(cv2.imread(seg_path))
predflow_var = predflow.unsqueeze(0)
gtflow_var = gtflow.unsqueeze(0)
segmask_var = segmask.unsqueeze(0)
predflow_var = predflow_var.to(device)
gtflow_var = gtflow_var.to(device)
segmask_var = segmask_var.to(device)
# compute output
epe = realEPE(predflow_var, gtflow_var)
epe_parts = partsEPE(predflow_var, gtflow_var, segmask_var)
#epe_parts.update((x, args.div_flow*y) for x, y in epe_parts.items() )
# record EPE
flow_epe.update(epe.item(), gtflow_var.size(0))
for bk in avg_parts_epe:
if epe_parts[bk].item() > 0:
avg_parts_epe[bk].update(epe_parts[bk].item(),
gtflow_var.size(0))
epe_dict = {}
for bk in BODY_MAP.keys():
epe_dict[bk] = avg_parts_epe[bk].avg
epe_dict['full_epe'] = flow_epe.avg
np.save(os.path.join('results', args.save_name), epe_dict)
print("Averge EPE", flow_epe.avg)
def main():
global args, best_EPE, save_path
args = parser.parse_args()
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
save_path = data_dir / 'flow'
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_pairs = []
for ext in args.img_exts:
test_files = data_dir.files('*0.{}'.format(ext))
for file in test_files:
img_pair = file.parent / (file.namebase[:-1] + '1.{}'.format(ext))
if img_pair.isfile():
img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).cuda()
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = input_transform(imread(img1_file))
img2 = input_transform(imread(img2_file))
input_var = torch.autograd.Variable(torch.cat([img1, img2], 0).cuda(),
volatile=True).unsqueeze(0)
# compute output
output = model(input_var)
rgb_flow = flow2rgb(args.div_flow * output.data[0].cpu().numpy(),
max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8)
imsave(save_path / (img1_file.namebase[:-2] + '_flow.png'), to_save)
def run(img1, img2, bidirectional=False, upsampling='bilinear', div_flow=20, max_flow=None, to_rgb=False):
if model is None:
setup()
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255]),
transforms.Normalize(mean=[0.411,0.432,0.45], std=[1,1,1])
])
img1 = input_transform(np.array(img1))
img2 = input_transform(np.array(img2))
input_var = torch.cat([img1, img2]).unsqueeze(0)
if bidirectional:
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
flow = model(input_var)
if upsampling is not None:
assert upsampling in ['nearest', 'bilinear'], \
'Upsampling mode {} not recognized'.format(upsampling)
flow = torch.nn.functional.interpolate(
flow,
size=img1.size()[-2:],
mode=upsampling,
align_corners=False)
if to_rgb:
rgb_flow = flow2rgb(div_flow * flow[0], max_value=max_flow)
rgb_flow = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0)
return rgb_flow
else:
flow = (div_flow * flow[0]).detach().cpu().numpy().transpose(1,2,0)
return flow
def main():
global args
args = parser.parse_args()
test_list = make_dataset(args.data)
# test_list = make_real_dataset(args.data)
if args.arch == 'pwc':
model = models.pwc_dc_net('models/pwc_net_ft.pth.tar').cuda()
elif args.arch == 'spynet':
model = models.spynet(nlevels=5, strmodel='F').cuda()
elif args.arch == 'flownet2':
model = models.FlowNet2().cuda()
print("=> using pre-trained weights for FlowNet2")
weights = torch.load('models/FlowNet2_checkpoint.pth.tar')
model.load_state_dict(weights['state_dict'])
if args.pretrained is not None:
network_data = torch.load(args.pretrained)
args.arch = network_data['arch']
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](data=network_data).cuda()
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
model.eval()
flow_epe = AverageMeter()
avg_mot_err = AverageMeter()
avg_parts_epe = {}
for bk in BODY_MAP.keys():
avg_parts_epe[bk] = AverageMeter()
if args.no_norm:
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255])
])
else:
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0], std=[args.div_flow, args.div_flow])
])
for i, (img_paths, flow_path, seg_path) in enumerate(tqdm(test_list)):
raw_im1 = flow_transforms.ArrayToTensor()(imread(img_paths[0],
mode='RGB')[:, :, :3])
raw_im2 = flow_transforms.ArrayToTensor()(imread(img_paths[1],
mode='RGB')[:, :, :3])
img1 = input_transform(imread(img_paths[0], mode='RGB')[:, :, :3])
img2 = input_transform(imread(img_paths[1], mode='RGB')[:, :, :3])
if flow_path is None:
_, h, w = img1.size()
new_h = int(np.floor(h / 256) * 256)
new_w = int(np.floor(w / 448) * 448)
# if i>744:
# import ipdb; ipdb.set_trace()
img1 = F.upsample(img1.unsqueeze(0), (new_h, new_w),
mode='bilinear').squeeze()
img2 = F.upsample(img2.unsqueeze(0), (new_h, new_w),
mode='bilinear').squeeze()
if flow_path is not None:
gtflow = target_transform(load_flo(flow_path))
segmask = flow_transforms.ArrayToTensor()(cv2.imread(seg_path))
input_var = torch.cat([img1, img2]).unsqueeze(0)
if flow_path is not None:
gtflow_var = gtflow.unsqueeze(0)
segmask_var = segmask.unsqueeze(0)
input_var = input_var.to(device)
if flow_path is not None:
gtflow_var = gtflow_var.to(device)
segmask_var = segmask_var.to(device)
# compute output
output = model(input_var)
if flow_path is not None:
epe = args.div_flow * realEPE(
output,
gtflow_var,
sparse=True if 'KITTI' in args.dataset else False)
epe_parts = partsEPE(output, gtflow_var, segmask_var)
epe_parts.update(
(x, args.div_flow * y) for x, y in epe_parts.items())
# record EPE
flow_epe.update(epe.item(), gtflow_var.size(0))
for bk in avg_parts_epe:
if epe_parts[bk].item() > 0:
avg_parts_epe[bk].update(epe_parts[bk].item(),
gtflow_var.size(0))
# record motion warping error
raw_im1 = raw_im1.cuda().unsqueeze(0)
raw_im2 = raw_im2.cuda().unsqueeze(0)
mot_err = motion_warping_error(raw_im1, raw_im2,
args.div_flow * output)
avg_mot_err.update(mot_err.item(), raw_im1.size(0))
if args.output_dir is not None:
if flow_path is not None:
_, h, w = gtflow.size()
output_path = flow_path.replace(args.data, args.output_dir)
output_path = output_path.replace('/test/', '/')
os.system('mkdir -p ' + output_path[:-15])
else:
output_path = img_paths[0].replace(args.data, args.output_dir)
os.system('mkdir -p ' + output_path[:-10])
output_path = output_path.replace('.png', '.flo')
output_path = output_path.replace('/flow/', '/')
upsampled_output = F.interpolate(output, (h, w),
mode='bilinear',
align_corners=False)
flow_write(output_path,
upsampled_output.cpu()[0].data.numpy()[0],
upsampled_output.cpu()[0].data.numpy()[1])
if args.save_name is not None:
epe_dict = {}
for bk in BODY_MAP.keys():
epe_dict[bk] = avg_parts_epe[bk].avg
epe_dict['full_epe'] = flow_epe.avg
np.save(os.path.join('results', args.save_name), epe_dict)
print("Averge EPE", flow_epe.avg)
print("Motion warping error", avg_mot_err.avg)
def main():
global args, save_path
args = parser.parse_args()
data_dir = args.data
print("=> fetching img pairs in '{}'".format(data_dir))
if args.output is None:
save_path = args.data / 'flow_' + subset
else:
save_path = args.output # Path(args.output)
print('=> will save everything to {}'.format(save_path))
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
# create model
network_data = torch.load(args.pretrained, map_location=device)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
#
# Use encoder (from the first conv layer to conv6_1) of FlowNetS as a feature extractor # added by Thao
for param in model.parameters():
param.requires_grad = False
modules = list(model.children())[0:9]
feature_extractor = nn.Sequential(*modules)
#
cudnn.benchmark = True
# List of RGB images
img_files = []
csv_file = 'frames.csv' if extension == "" else 'extension.csv'
with open(os.path.join(args.csvpath, csv_file), 'r') as csvfile:
csvReader = csv.reader(csvfile)
for row in csvReader:
temp = row[0]
img_files.append(temp)
num_files = len(img_files)
num_stacks = int(num_files / RGB_stack_size)
print('Num_stacks: ', num_stacks)
start_time = time.time()
if num_files < 2:
print('Check the image folder')
else:
all_features = []
num_movie = 0
for stackID in range(0, num_stacks):
clip_all_features = []
count = 0
start_time = time.time()
for k in range(int(stackID * RGB_stack_size),
int((stackID + 1) * RGB_stack_size) - 1):
im1_fn = img_files[k]
im2_fn = img_files[k + 1]
# Extract OF features
img1 = input_transform(imread(os.path.join(data_dir, im1_fn)))
img2 = input_transform(imread(os.path.join(data_dir, im2_fn)))
input_var = torch.cat([img1, img2]).unsqueeze(0)
#
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
#
input_var = input_var.to(device)
# compute feature map
output = feature_extractor(input_var)
# output is a feature map of [1, 1024, 6 ,8] => apply AvgPooling2D to get [1, 1024]-feature vector
create_AP2d = nn.AvgPool2d((output.size(2), output.size(3)),
stride=None)
feature_vector = create_AP2d(output).view(
output.size(0),
output.size(1)) # feature_vector is a [1, 1024] vector
# from tensor to numpy
feature_vector = np.float32(
feature_vector.cpu().detach().numpy())
#
clip_all_features.append(feature_vector)
#
#im1_fn = im2_fn
#
count += 1
if (k == int((stackID + 1) * RGB_stack_size) - 2):
clip_all_features.append(feature_vector)
num_movie += 1
print('Processed', num_movie, ' movie excerpts')
print('Running time for a movie excerpt: ',
time.time() - start_time)
start_time = time.time()
break
temp = np.stack(clip_all_features, axis=0)
temp = temp.squeeze(1)
avg_clip_features_8_seg = np.mean(temp, axis=0)
# Create a list of feature vectors
all_features.append(avg_clip_features_8_seg)
all_features = np.stack(all_features, axis=0)
print('Number of feature vectors: ', all_features.shape)
# save i3d_features in a .h5 file
start_time = time.time()
filename = os.path.join(args.output,
'FlowNetS_features' + extension + '.h5')
h5file = h5py.File(filename, mode='w')
h5file.create_dataset('data',
data=np.array(all_features, dtype=np.float32))
h5file.close()
print("Time for writing feature vectors in .h5 file: ",
time.time() - start_time)
def main():
global args, best_EPE
args = parser.parse_args()
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.45, 0.432, 0.411], std=[1, 1, 1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0], std=[args.div_flow, args.div_flow])
])
if 'KITTI' in args.dataset:
args.sparse = True
if args.sparse:
co_transform = flow_transforms.Compose([
flow_transforms.RandomCrop((320, 448)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
else:
co_transform = flow_transforms.Compose([
flow_transforms.RandomTranslate(10),
flow_transforms.RandomRotate(10, 5),
flow_transforms.RandomCrop((320, 448)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
print("=> fetching img pairs in '{}'".format(args.data))
train_set, test_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
target_transform=target_transform,
co_transform=co_transform,
split=args.split_file if args.split_file else args.split_value)
print('{} samp-les found, {} train samples and {} test samples '.format(
len(test_set) + len(train_set), len(train_set), len(test_set)))
# train_loader = torch.utils.data.DataLoader(
# train_set, batch_size=args.batch_size,
# num_workers=args.workers, pin_memory=True, shuffle=True)
val_loader = torch.utils.data.DataLoader(test_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=False)
# create model
if args.pretrained:
if torch.cuda.is_available():
network_data = torch.load(args.pretrained)
else:
network_data = torch.load(args.pretrained,
map_location=torch.device('cpu'))
# args.arch = network_data['arch']
print("=> using pre-trained model '{}'".format(args.arch))
else:
network_data = None
print("=> creating model '{}'".format(args.arch))
if args.qw and args.qa is not None:
if torch.cuda.is_available():
model = models.__dict__[args.arch](data=network_data,
bitW=args.qw,
bitA=args.qa).cuda()
else:
model = models.__dict__[args.arch](data=network_data,
bitW=args.qw,
bitA=args.qa)
else:
if torch.cuda.is_available():
model = models.__dict__[args.arch](data=network_data).cuda()
else:
model = models.__dict__[args.arch](data=network_data)
# if torch.cuda.is_available():
# model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
assert (args.solver in ['adam', 'sgd'])
print('=> setting {} solver'.format(args.solver))
param_groups = [{
'params': model.bias_parameters(),
'weight_decay': args.bias_decay
}, {
'params': model.weight_parameters(),
'weight_decay': args.weight_decay
}]
if args.solver == 'adam':
optimizer = torch.optim.Adam(param_groups,
args.lr,
betas=(args.momentum, args.beta))
elif args.solver == 'sgd':
optimizer = torch.optim.SGD(param_groups,
args.lr,
momentum=args.momentum)
# print (summary(model, (6, 320, 448)))
# print (model)
if args.evaluate:
best_EPE = validate(val_loader, model, 0)
# validate_chairs(model.module)
# validate_sintel(model.module)
return
def main(args):
save_path = 'epoch={},batchsize={},lr={},lambda={},lin_reso={},non_resolu={}'.format(
args.epochs, args.batch_size, args.lr, args.lambda_list,
args.linear_resolution, args.nonlinear_resolution)
print('=> will save everything to {}'.format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0], std=[255]),
transforms.Normalize(mean=[0.233], std=[0.255]),
])
print("=> fetching img pairs in '{}'".format(args.data))
train_set, test_set = cremidataset(args.data,
transform=input_transform,
split=0.9)
print('{} samples found, {} train samples and {} test samples '.format(
len(test_set) + len(train_set), len(train_set), len(test_set)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
num_workers=min(
args.batch_size, 8),
pin_memory=True,
shuffle=True)
val_loader = torch.utils.data.DataLoader(test_set,
batch_size=1,
num_workers=1,
pin_memory=True,
shuffle=False)
model = Baseline(18, 18, args).cuda()
cudnn.benchmark = True
optimizer = torch.optim.Adam(model.parameters(),
args.lr,
betas=(0.9, 0.999))
lr_para = filter(lambda p: p.requires_grad, model.parameters())
N = sum([numpy.prod(p.size()) for p in lr_para])
print(N)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.milestones, gamma=0.5)
best_EPE = -1
for epoch in range(args.epochs):
train_loss = train(train_loader, model, optimizer, epoch)
print(train_loss)
with torch.no_grad():
test_loss = validate(val_loader, model)
scheduler.step()
if best_EPE < 0:
best_EPE = test_loss
is_best = test_loss < best_EPE
best_EPE = min(test_loss, best_EPE)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
}, is_best, save_path)
def main():
global args, save_path
args = parser.parse_args()
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
save_path = data_dir / 'flow'
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_pairs = []
for ext in args.img_exts:
test_files = data_dir.files('*0.{}'.format(ext))
for file in test_files:
img_pair = file.parent / (file.namebase[:-1] + '1.{}'.format(ext))
if img_pair.isfile():
img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).cuda()
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
input_transform = test_transform()
for (img1_file, img2_file) in tqdm(img_pairs):
# img1 = input_transform(imread(img1_file))
# img2 = input_transform(imread(img2_file))
img1 = Image.open(img1_file).convert('RGB')
img1 = input_transform(img1)
img2 = Image.open(img2_file).convert('RGB')
img2 = input_transform(img2)
input_var = torch.tensor(torch.cat([img1, img2]).cuda()).unsqueeze(0)
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.tensor(
torch.cat([img2, img1], 0).cuda()).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
# compute output
output = model(input_var)
args.upsampling = 'bilinear'
if args.upsampling is not None:
output = torch.nn.functional.upsample(
output, size=img1.size()[-2:], mode=args.upsampling) * 80
# output = output.permute(0, 2, 3, 1) * 4
# img1 = torch.nn.functional.avg_pool2d(img1, 4)
# img2 = torch.nn.functional.avg_pool2d(img2, 4)
# N, C, H, W = output.size()
print(torch.sum(output > 3), flush=True)
from PWCNet import warp
'''
output = torch.zeros(N, C, H, W)
output.fill_(10.0)
output = output.cuda()
'''
result = warp(img2.unsqueeze(0).cuda(), 1.0 * output)
print(result.size())
save_image(result, 'img2to.jpg')
save_image(img1, 'img1.jpg')
save_image(img2, 'img2.jpg')
for suffix, flow_output in zip(['flow', 'inv_flow'],
output.data.cpu()):
rgb_flow = flow2rgb(args.div_flow * flow_output,
max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imsave(
save_path / '{}{}.png'.format(img1_file.namebase[:-1], suffix),
to_save)
def main():
global args, best_EPE, save_path, intrinsic
# ============= savor setting ===================
save_path = '{}_{}_{}epochs{}_b{}_lr{}_posW{}'.format(
args.arch,
args.solver,
args.epochs,
'_epochSize'+str(args.epoch_size) if args.epoch_size > 0 else '',
args.batch_size,
args.lr,
args.pos_weight,
)
if not args.no_date:
timestamp = datetime.datetime.now().strftime("%y_%m_%d_%H_%M")
else:
timestamp = ''
save_path = os.path.abspath(args.savepath) + '/' + os.path.join(args.dataset, save_path + '_' + timestamp )
# ========== Data loading code ==============
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255]),
transforms.Normalize(mean=[0.411,0.432,0.45], std=[1,1,1])
])
val_input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
])
co_transform = flow_transforms.Compose([
flow_transforms.RandomCrop((args.train_img_height ,args.train_img_width)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
print("=> loading img pairs from '{}'".format(args.data))
train_set, val_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
val_transform = val_input_transform,
target_transform=target_transform,
co_transform=co_transform
)
print('{} samples found, {} train samples and {} val samples '.format(len(val_set)+len(train_set),
len(train_set),
len(val_set)))
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=args.batch_size,
num_workers=args.workers, pin_memory=True, shuffle=True, drop_last=True)
val_loader = torch.utils.data.DataLoader(
val_set, batch_size=args.batch_size,
num_workers=args.workers, pin_memory=True, shuffle=False, drop_last=True)
# ============== create model ====================
if args.pretrained:
network_data = torch.load(args.pretrained)
args.arch = network_data['arch']
print("=> using pre-trained model '{}'".format(args.arch))
else:
network_data = None
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]( data = network_data).cuda()
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
#=========== creat optimizer, we use adam by default ==================
assert(args.solver in ['adam', 'sgd'])
print('=> setting {} solver'.format(args.solver))
param_groups = [{'params': model.module.bias_parameters(), 'weight_decay': args.bias_decay},
{'params': model.module.weight_parameters(), 'weight_decay': args.weight_decay}]
if args.solver == 'adam':
optimizer = torch.optim.Adam(param_groups, args.lr,
betas=(args.momentum, args.beta))
elif args.solver == 'sgd':
optimizer = torch.optim.SGD(param_groups, args.lr,
momentum=args.momentum)
# for continues training
if args.pretrained and ('dataset' in network_data):
if args.pretrained and args.dataset == network_data['dataset'] :
optimizer.load_state_dict(network_data['optimizer'])
best_EPE = network_data['best_EPE']
args.start_epoch = network_data['epoch']
save_path = os.path.dirname(args.pretrained)
print('=> will save everything to {}'.format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
train_writer = SummaryWriter(os.path.join(save_path, 'train'))
val_writer = SummaryWriter(os.path.join(save_path, 'val'))
# spixelID: superpixel ID for visualization,
# XY_feat: the coordinate feature for position loss term
spixelID, XY_feat_stack = init_spixel_grid(args)
val_spixelID, val_XY_feat_stack = init_spixel_grid(args, b_train=False)
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
train_avg_slic, train_avg_sem, iteration = train(train_loader, model, optimizer, epoch,
train_writer, spixelID, XY_feat_stack )
if epoch % args.record_freq == 0:
train_writer.add_scalar('Mean avg_slic', train_avg_slic, epoch)
# evaluate on validation set and save the module( and choose the best)
with torch.no_grad():
avg_slic, avg_sem = validate(val_loader, model, epoch, val_writer, val_spixelID, val_XY_feat_stack)
if epoch % args.record_freq == 0:
val_writer.add_scalar('Mean avg_slic', avg_slic, epoch)
rec_dict = {
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.module.state_dict(),
'best_EPE': best_EPE,
'optimizer': optimizer.state_dict(),
'dataset': args.dataset
}
if (iteration) >= (args.milestones[-1] + args.additional_step):
save_checkpoint(rec_dict, is_best =False, filename='%d_step.tar' % iteration)
print("Train finished!")
break
if best_EPE < 0:
best_EPE = avg_sem
is_best = avg_sem < best_EPE
best_EPE = min(avg_sem, best_EPE)
save_checkpoint(rec_dict, is_best)
def main():
torch.cuda.set_device(1)
global args, best_EPE, save_path
args = parser.parse_args()
save_path = '{},{},{}epochs{},b{},lr{}'.format(
args.arch,
args.solver,
args.epochs,
',epochSize'+str(args.epoch_size) if args.epoch_size > 0 else '',
args.batch_size,
args.lr)
if not args.no_date:
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
save_path = os.path.join(timestamp,save_path)
save_path = os.path.join(args.dataset,save_path)
print('=> will save everything to {}'.format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
###tensorboard
train_writer = SummaryWriter(os.path.join(save_path,'train'))
test_writer = SummaryWriter(os.path.join(save_path,'test'))
training_output_writers = []
for i in range(15):
training_output_writers.append(SummaryWriter(os.path.join(save_path,'train',str(i))))
output_writers = []
for i in range(15):
output_writers.append(SummaryWriter(os.path.join(save_path,'test',str(i))))
################
######## Data transformation code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(), ##from numpy array to tensor
transforms.Normalize(mean=[0,0,0], std=[255,255,255]) ##divide each channel of the image by 255
]) ###input image transform
background_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255])
]) ####background image transform
gradient_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255])
]) ##gradient transform
reflection_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255])
]) ##reflection transform
co_transform = flow_transforms.Compose([
flow_transforms.RandomVerticalFlip(), ##flip the image vertically
flow_transforms.RandomHorizontalFlip(), ##flip the image horizontally
#flow_transforms.RandomColorWarp(0,0)
])
#####data loading see the mpi_sintel_both file in datasets folder
print("=> fetching img pairs in '{}'".format(args.data))
train_set, test_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
gradient_transform = gradient_transform,
reflection_transform = reflection_transform,
background_transform = background_transform,
co_transform = co_transform
)
print('{} samples found, {} train samples, {} test samples '.format(len(train_set), len(train_set), len(test_set)))
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=args.batch_size,
num_workers=0, pin_memory=True, shuffle=True)
val_loader = torch.utils.data.DataLoader(
test_set, batch_size=args.batch_size,
num_workers=0, pin_memory=True, shuffle=False)
vgg = torchvision.models.vgg16_bn(pretrained = True) ##load the vgg model
vgglist = list(vgg.features.children())
model = ReflectionNetwork(vgglist) ##load the training model
optimizer = optim.Adam(model.parameters(), lr = args.lr)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=0.5)
model = model.cuda()
if args.pretrained:
data = torch.load('/home/Test25/model2/checkpoint_81.pth.tar')
model.load_state_dict(data['state_dict'])
####load the loss functions
ssimLoss = SSIMLoss().cuda()
L1Loss = nn.L1Loss().cuda()
siLoss = SILoss().cuda()
mmdLoss = MMDLoss().cuda()
num_epoches = 80
for epoch in range(num_epoches+1):
scheduler.step()
print('epoch {}'.format(epoch))
train(train_loader, optimizer, ssimLoss, L1Loss, siLoss, mmdLoss, model, train_writer, training_output_writers, epoch) ##training code
validate(val_loader, model, L1Loss, test_writer, output_writers, epoch)
if epoch % 1 == 0:
#save the model
save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer' : optimizer.state_dict(),
},epoch+1)
def main():
global args, best_EPE, save_path
args = parser.parse_args()
save_path = '{},{},{}epochs{},b{},lr{}'.format(
args.arch, args.solver, args.epochs,
',epochSize' + str(args.epoch_size) if args.epoch_size > 0 else '',
args.batch_size, args.lr)
if not args.no_date:
timestamp = datetime.datetime.now().strftime("%a-%b-%d-%H:%M")
save_path = os.path.join(timestamp, save_path)
save_path = os.path.join(args.dataset, save_path)
print('=> will save everything to {}'.format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
# Data loading code
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]), normalize
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0], std=[args.div_flow, args.div_flow])
])
if 'KITTI' in args.dataset:
co_transform = flow_transforms.Compose([
flow_transforms.RandomCrop((320, 448)),
#random flips are not supported yet for tensor conversion, but will be
#flow_transforms.RandomVerticalFlip(),
#flow_transforms.RandomHorizontalFlip()
])
else:
co_transform = flow_transforms.Compose([
flow_transforms.RandomTranslate(10),
flow_transforms.RandomRotate(10, 5),
flow_transforms.RandomCrop((320, 448)),
#random flips are not supported yet for tensor conversion, but will be
#flow_transforms.RandomVerticalFlip(),
#flow_transforms.RandomHorizontalFlip()
])
print("=> fetching img pairs in '{}'".format(args.data))
train_set, test_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
target_transform=target_transform,
co_transform=co_transform,
split=args.split)
print('{} samples found, {} train samples and {} test samples '.format(
len(test_set) + len(train_set), len(train_set), len(test_set)))
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=args.batch_size,
sampler=balancedsampler.RandomBalancedSampler(train_set,
args.epoch_size),
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(test_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch](args.pretrained).cuda()
model = torch.nn.DataParallel(model).cuda()
criterion = multiscaleloss(sparse='KITTI' in args.dataset,
loss=args.loss).cuda()
high_res_EPE = multiscaleloss(scales=1,
downscale=4,
weights=(1),
loss='L1',
sparse='KITTI' in args.dataset).cuda()
cudnn.benchmark = True
assert (args.solver in ['adam', 'sgd'])
print('=> setting {} solver'.format(args.solver))
if args.solver == 'adam':
optimizer = torch.optim.Adam(model.parameters(),
args.lr,
betas=(args.momentum, args.beta),
weight_decay=args.weight_decay)
elif args.solver == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.evaluate:
best_EPE = validate(val_loader, model, criterion, high_res_EPE)
return
with open(os.path.join(save_path, args.log_summary), 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'train_EPE', 'EPE'])
with open(os.path.join(save_path, args.log_full), 'w') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow(['train_loss', 'train_EPE'])
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train_loss, train_EPE = train(train_loader, model, criterion,
high_res_EPE, optimizer, epoch)
# evaluate o validation set
EPE = validate(val_loader, model, criterion, high_res_EPE)
if best_EPE < 0:
best_EPE = EPE
# remember best prec@1 and save checkpoint
is_best = EPE < best_EPE
best_EPE = min(EPE, best_EPE)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.module.state_dict(),
'best_EPE': best_EPE,
}, is_best)
with open(os.path.join(save_path, args.log_summary), 'a') as csvfile:
writer = csv.writer(csvfile, delimiter='\t')
writer.writerow([train_loss, train_EPE, EPE])
NUM_VAL = 300
NUM_TEST = 149
listing = random.sample(os.listdir(input_rgb_images_dir), 1449)
train_listing = listing[:min(NUM_TRAIN, train_on)]
val_listing = listing[NUM_TRAIN:min(NUM_VAL + NUM_TRAIN, val_on + NUM_TRAIN)]
test_listing = listing[NUM_VAL +
NUM_TRAIN:min(NUM_VAL + NUM_TRAIN + NUM_TEST, test_on +
NUM_VAL + NUM_TRAIN)]
data_dir = (input_rgb_images_dir, target_depth_images_dir,
target_labels_images_dir)
input_transform = transforms.Compose(
[flow_transforms.Scale(228),
flow_transforms.ArrayToTensor()])
target_depth_transform = transforms.Compose(
[flow_transforms.Scale_Single(228),
flow_transforms.ArrayToTensor()])
target_labels_transform = transforms.Compose([flow_transforms.ArrayToTensor()])
##Apply this transform on input, ground truth depth images and labeled images
co_transform = flow_transforms.Compose([
flow_transforms.RandomCrop((480, 640)),
flow_transforms.RandomHorizontalFlip()
])
##Splitting in train, val and test sets [No data augmentation on val and test, only on train]
train_dataset = ListDataset(data_dir,train_listing,input_transform,target_depth_transform,\
def generate_flow(self):
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
# create model
network_data = torch.load(self.model_dir, map_location='cpu')
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).cuda()
model.eval()
cudnn.benchmark = True
# if 'div_flow' in network_data.keys():
# args.div_flow = network_data['div_flow']
mp4_list = glob.glob(os.path.join(self.args.input_dir, "*.mp4"))
print(self.args.input_dir)
print(mp4_list)
for mp4_fn in mp4_list:
print(mp4_fn)
mp4_ims = mp4_load(mp4_fn)
n_im = len(mp4_ims) # Number of images
export_ims = []
[nY, nX, nC] = mp4_ims[0].shape
outs = []
#frame_ints = list(range(n_im))
for ii in tqdm(range(n_im - args.window)):
aa = ii
bb = ii + args.window
img1 = input_transform(mp4_ims[aa])
img2 = input_transform(mp4_ims[bb])
input_var = torch.cat([img1, img2]).unsqueeze(0)
# if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.cuda()
# compute output
output = model(input_var)
outs.append(np.array(output.cpu().detach()))
# Upsample
output = F.interpolate(output,
size=img1.size()[-2:],
mode='bilinear',
align_corners=False)
# Convert to an RGBim1.
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
rgb_flow = flow2rgb(args.div_flow * flow_output,
max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(
1, 2, 0)
combined = np.zeros([nY, args.join * nX, nC], dtype=np.uint8)
if args.join == 2:
combined[:, :nX, :] = mp4_ims[aa][:, :, [
2, 1, 0
]] # Invert back for BGR
combined[:, nX:, :] = to_save
else:
combined[:, :, :] = to_save
export_ims.append(combined) # Save for list
self.save_flow(mp4_fn, export_ims, mp4_ims[0].shape)
def main():
#0. assert args
global args, save_path
args = parser.parse_args()
if args.output_value == 'both':
output_string = "raw output and RGB visualization"
elif args.output_value == 'raw':
output_string = "raw output"
elif args.output_value == 'vis':
output_string = "RGB visualization"
print("=> will save " + output_string)
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
if args.output is None:
save_path = data_dir / 'flow'
else:
save_path = Path(args.output)
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
#1. Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_pairs = []
for ext in args.img_exts:
test_files = data_dir.files('*1.{}'.format(ext))
for file in test_files:
img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
if img_pair.isfile():
img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
#2. create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(
device
) #models = {module}<module 'models' from '/home/roit/ws/flownet_pt/models/__init__.py'>
model.eval() #flownetc or flownets
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = imread(img1_file)
img2 = imread(img2_file)
img1 = input_transform(img1)[:3, :, :]
img2 = input_transform(img2)[:3, :, :]
input_var = torch.cat([img1, img2
]).unsqueeze(0) #input_var={Tensor}, [1,6,h,w]
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(
0) # 如果用flownetc 到时候再分开
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var) #outTensor [1,2,h/4,w/4]
if args.upsampling is not None:
output = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
out1 = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
filename = save_path / '{}{}'.format(img1_file.namebase[:-1],
suffix)
if args.output_value in ['vis', 'both']:
tmp = args.div_flow * flow_output
rgb_flow = flow2rgb(tmp, max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imwrite(filename + '.png', to_save)
if args.output_value in ['raw', 'both']:
# Make the flow map a HxWx2 array as in .flo files
to_save = (args.div_flow *
flow_output).cpu().numpy().transpose(1, 2, 0)
np.save(filename + '.npy', to_save)
def main():
global args, save_path
args = parser.parse_args()
if args.output_value == 'both':
output_string = "raw output and RGB visualization"
elif args.output_value == 'raw':
output_string = "raw output"
elif args.output_value == 'vis':
output_string = "RGB visualization"
print("=> will save " + output_string)
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
if args.output is None:
save_path = data_dir / 'flow'
else:
save_path = Path(args.output)
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
# Make sure the folder has extracted images in the format of "name + num + extension", where num start from 0, 1, 2, ...
# Currently only test on MPI-Sintel, which comprises .png images
img_pairs = []
img_list = os.listdir(data_dir)
img_list = sorted(img_list)
#img_list.sort()
img_list = img_list[1:]
print(img_list)
#for ext in args.img_exts:
for ind in range(len(img_list) - 1):
img1 = img_list[ind]
print("img1: ", img1)
#test_files = data_dir.files('*1.{}'.format(ext))
img1_name, ext = img1.split('.')
basename, num = img1_name.split('_')
img2 = '.'.join(
['_'.join([basename, str(int(num) + 1).zfill(4)]), ext])
print("img 2: ", img2)
img_pairs.append([data_dir + '/' + img1, data_dir + '/' + img2])
#print("Test files: ")
#print(test_files)
#for file in test_files:
#img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
#print("file: ")
#print(file)
#print("img_pair: ")
#print(img_pair)
#if img_pair.isfile():
# img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
#print(img_pairs)
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
output = data_dir + '/flow/' + 'output.mp4'
print(output)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output, fourcc, 10.0, (640, 480))
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = input_transform(imread(img1_file))
img2 = input_transform(imread(img2_file))
input_var = torch.cat([img1, img2]).unsqueeze(0)
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var)
if args.upsampling is not None:
output = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
print(img1_file.namebase)
filename = save_path / '{}{}'.format(img1_file.namebase, suffix)
print(filename)
if args.output_value in ['vis', 'both']:
rgb_flow = flow2rgb(args.div_flow * flow_output,
max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imwrite(filename + '.png', to_save)
out.write(to_save)
cv2.imshow('video', to_save)
if (cv2.waitKey(1) & 0xFF) == ord('q'): break
if args.output_value in ['raw', 'both']:
# Make the flow map a HxWx2 array as in .flo files
to_save = (args.div_flow *
flow_output).cpu().numpy().transpose(1, 2, 0)
np.save(filename + '.npy', to_save)
out.release()
cv2.destroyAllWindows()
def main():
global args, save_path
args = parser.parse_args()
data_dir = args.data
print("=> fetching img pairs in '{}'".format(args.data))
save_path = args.output
print('=> will save everything to {}'.format(save_path))
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_pairs = []
names = []
# for ext in args.img_exts:
# test_files = data_dir.files('*1.{}'.format(ext))
# for file in test_files:
# img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
# if img_pair.isfile():
# img_pairs.append([file, img_pair])
# print(img_pairs)
canon_im1num = [
1, 1, 2, 10, 10, 11, 21, 21, 21, 21, 22, 22, 22, 24, 25, 31, 31, 31,
32, 32, 33, 41, 43, 51, 54
] # 25
canon_im2num = [
2, 3, 3, 11, 12, 12, 22, 24, 25, 26, 23, 24, 27, 26, 27, 32, 33, 34,
33, 34, 34, 42, 45, 52, 55
]
# im1 = '%s/1/(1).JPG'%data_dir
# im2 = '%s/2/(1).JPG'%data_dir
# img_pairs.append([im1,im2])
s = 0
for i in range(0, 25):
foldernum1 = canon_im1num[i]
foldernum2 = canon_im2num[i]
rawnum = len(glob.glob('../../data/canon/%d/*.CR2' % (foldernum1)))
for j in range(rawnum):
image_path1 = '../../data/canon/%d/(%d).CR2' % (foldernum1, j + 1)
image_path2 = '../../data/canon/%d/(%d).CR2' % (foldernum2, j + 1)
img_pairs.append([image_path1, image_path2])
n = '%d-%d_%d' % (foldernum1, foldernum2, j + 1)
names.append(n)
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
count = -1
for (img1_file, img2_file) in tqdm(img_pairs):
raw1 = rawpy.imread(img1_file)
raw2 = rawpy.imread(img2_file)
im1 = raw1.raw_image_visible.astype(np.float32)
im1 = (im1 - 2047) / (16383 - 2047)
ratio = 0.3 / np.mean(im1)
im1 = np.minimum(np.maximum(im1 * ratio, 0.0), 1.0)
im2 = raw2.raw_image_visible.astype(np.float32)
im2 = (im2 - 2047) / (16383 - 2047)
ratio = 0.3 / np.mean(im2)
im2 = np.minimum(np.maximum(im2 * ratio, 0.0), 1.0)
# print(ratio)
im1 = np.expand_dims(im1, axis=2)
H = im1.shape[0]
W = im1.shape[1]
image1 = np.concatenate(
(
im1[0:H:2, 0:W:2, :], #r
(im1[0:H:2, 1:W:2, :] + im1[1:H:2, 0:W:2, :]) / 2.0, #g
im1[1:H:2, 1:W:2, :]),
axis=2) #b
im2 = np.expand_dims(im2, axis=2)
H = im2.shape[0]
W = im2.shape[1]
image2 = np.concatenate(
(
im2[0:H:2, 0:W:2, :], #r
(im2[0:H:2, 1:W:2, :] + im2[1:H:2, 0:W:2, :]) / 2.0, #g
im2[1:H:2, 1:W:2, :]),
axis=2) #b
i1 = crop(image1, 1920, 2944, 2)
i2 = crop(image2, 1920, 2944, 2)
i1 = (i1 * 255).astype('uint8')
i2 = (i2 * 255).astype('uint8')
i1 = cv.fastNlMeansDenoisingColored(i1, None, 10, 10, 7, 21)
i2 = cv.fastNlMeansDenoisingColored(i2, None, 10, 10, 7, 21)
img1 = input_transform(i1)
img2 = input_transform(i2)
# print(i1.shape)
# print(img1.shape)
input_var = torch.cat([img1, img2]).unsqueeze(0)
# if args.bidirectional:
# # feed inverted pair along with normal pair
# inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
# input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var)
if args.upsampling is not None:
output = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
count += 1
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
filename = '%s/%s' % (save_path, names[count])
print(filename)
# if args.output_value in['vis', 'both']:
rgb_flow = flow2rgb(flow_output, max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imwrite(filename + '.png', to_save)
# if args.output_value in ['raw', 'both']:
# # Make the flow map a HxWx2 array as in .flo files
to_save = (flow_output).cpu().numpy().transpose(1, 2, 0)
# print(to_save.shape)
flow_write(to_save, filename + '.flo')
def main():
global args, save_path
args = parser.parse_args()
save_path = args.output
print('=> will save everything to {}'.format(save_path))
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255]),
transforms.Normalize(mean=[0.411,0.432,0.45], std=[1,1,1])
])
img_pairs = []
names = []
# for ext in args.img_exts:
# test_files = data_dir.files('*1.{}'.format(ext))
# for file in test_files:
# img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
# if img_pair.isfile():
# img_pairs.append([file, img_pair])
# print(img_pairs)
canon_im1num = [1,1,2,10,10,11,21,21,21,21,22,22,22,24,25,31,31,31,32,32,33,41,43,51,54] # 25
canon_im2num = [2,3,3,11,12,12,22,24,25,26,23,24,27,26,27,32,33,34,33,34,34,42,45,52,55]
sony_im1num = [1,1,1,2,11,11,12,15] # 8
sony_im2num = [2,3,4,4,12,13,13,16]
# im1 = '%s/1/(1).JPG'%data_dir
# im2 = '%s/2/(1).JPG'%data_dir
# img_pairs.append([im1,im2])
s = 0
for i in range(0,8):
foldernum1 = sony_im1num[i]
foldernum2 = sony_im2num[i]
rawnum = len(glob.glob('../../data/sid_Sony/%d/*.png'%(foldernum1)))
for j in range(rawnum):
image_path1 = '../../data/nlm/%d_%d.jpg'%(foldernum1,j+1)
image_path2 = '../../data/nlm/%d_%d.jpg'%(foldernum2,j+1)
img_pairs.append([image_path1, image_path2])
n = '%d-%d_%d'%(foldernum1,foldernum2,j+1)
names.append(n)
# create model
network_data = torch.load(args.pretrained)
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
count = -1
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = input_transform(imread(img1_file))
img2 = input_transform(imread(img2_file))
# print(i1.shape)
# print(img1.shape)
input_var = torch.cat([img1, img2]).unsqueeze(0)
# if args.bidirectional:
# # feed inverted pair along with normal pair
# inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
# input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
# compute output
output = model(input_var)
if args.upsampling is not None:
output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False)
count += 1
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
filename = '%s/%s'%(save_path,names[count])
print(filename)
# if args.output_value in['vis', 'both']:
rgb_flow = flow2rgb(flow_output, max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0)
imwrite(filename + '.png', to_save)
# if args.output_value in ['raw', 'both']:
# # Make the flow map a HxWx2 array as in .flo files
to_save = (flow_output).cpu().numpy().transpose(1,2,0)
# print(to_save.shape)
flow_write(to_save, filename+'.flo')
def main():
global args, best_EPE
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
print(device)
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
# this takes too long
'''
datasets_path = "C:\data\splits"
test_data = np.load(os.path.join(datasets_path, "test_data.npy"))
test_gt = np.load(os.path.join(datasets_path, "test_gt.npy"))
train_data = np.load(os.path.join(datasets_path, "train_data.npy"))
train_gt = np.load(os.path.join(datasets_path, "train_gt.npy"))
print(np.shape(test_data))
'''
save_path = '..//data//output'
train_writer = SummaryWriter(os.path.join(save_path, 'train'))
test_writer = SummaryWriter(os.path.join(save_path, 'test'))
output_writers = []
for i in range(3):
output_writers.append(SummaryWriter(os.path.join(save_path, 'test', str(i))))
time_range = 64
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=dataset_mean, std=dataset_std),
#transforms.Normalize(mean=[0.45,0.432,0.411], std=[1,1,1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0],std=[args.div_flow,args.div_flow])
])
#co_transform = transforms.Compose([
# flow_transforms.RandomRotate(5),
# flow_transforms.RandomTranslate(10)
#])
#co_transform = flow_transforms.RandomRotate(10)
if RGB:
FM_1_testdir = "F://Flying_monkeys_2_RGB//test"
FM_1_traindir = "F://Flying_monkeys_2_RGB//train"
else:
FM_1_testdir = "C://data//FlyingMonkeys_1//test"
FM_1_traindir = "C://data//FlyingMonkeys_1//train"
train_data = []
train_gt = []
test_data = []
test_gt = []
for filename in os.listdir(FM_1_traindir):
if filename.startswith("data"):
#this is so hack-y...
train_data.append([filename, int(filename.split('_')[1].split('.')[0])]) #split the filename
# by the underscore and split the number from the filetype
elif filename.startswith("gt"):
train_gt.append([filename, int(filename.split('_')[1].split('.')[0])])
for filename in os.listdir(FM_1_testdir):
if filename.startswith("data"):
#this is so hack-y...
test_data.append([filename, int(filename.split('_')[1].split('.')[0])]) #split the filename
# by the underscore and split the number from the filetype
elif filename.startswith("gt"):
test_gt.append([filename, int(filename.split('_')[1].split('.')[0])])
train_data.sort(key = lambda x: x[1])
train_gt.sort(key=lambda x: x[1])
test_data.sort(key=lambda x: x[1])
test_gt.sort(key=lambda x: x[1])
train_list = [[data[0],gt[0]] for data, gt in zip(train_data,train_gt)]
test_list = [[data[0], gt[0]] for data, gt in zip(test_data, test_gt)]
#print(train_list [787])
#print(test_list[787])
#train_dataset = ListDataset(FM_1_traindir , train_list, input_transform, target_transform, co_transform)
#test_dataset = ListDataset(FM_1_testdir, test_list, input_transform, target_transform, co_transform)
train_dataset = ListDataset(FM_1_traindir, train_list, input_transform, target_transform)
test_dataset = ListDataset(FM_1_testdir, test_list, input_transform, target_transform)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size,
num_workers = 2, pin_memory=True, shuffle=True)
val_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=args.batch_size,
num_workers = 2, pin_memory=True, shuffle=False)
#a,b = train_dataset.__getitem__(178)
#print(np.shape(a))
#viewFlow(b,256,192)
#np.shape(a)
#print(a[: ,: ,0].dtype)
#print(a[:, :, 1].dtype)
#viewImg(a[:,:,1])
#viewImg(a[:, :, 0])
#print(a[100,100,0])
#print(a[100, 100, 1])
# Load model
model = PhotonNetS(batchNorm=True).to(device)
'''######################################################'''
# Evaluate on the test set using saved model
if evaluate:
assert os.path.exists(os.path.join(save_path, "datset1_15epochs.pt"))
# Set the test model
model = PhotonNetS(batchNorm=True).to(device)
model.load_state_dict(torch.load(os.path.join(save_path, "datset1_15epochs.pt")))
# test_loader = torch.utils.data.DataLoader(
# test_dataset, batch_size=args.test_batch_size, shuffle=True, **kwargs)
with torch.no_grad():
epoch = 1
EPE, total_photons = validate(val_loader, model, epoch, output_writers)
print("EPE is:", EPE)
return
if train:
# Try different optimzers here [Adam, SGD, RMSprop]
#optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
#optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.4)
optimizer = optim.Adam(model.parameters(), 0.00100, betas=(.9, .999))
# betas=(args.momentum, args.beta))
# Set learning rate scheduler
scheduler = StepLR(optimizer, step_size=args.step, gamma=args.gamma)
#scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=0.5)
'''
# Training loop
train_losses = []
test_losses = []
x = []
fig, ax = plt.subplots(1)
'''
train_EPE_list = []
test_EPE_list = []
epoch_list = []
for epoch in range(21):
epoch_list.append(epoch)
# train for one epoch
train_loss, train_EPE = train(train_loader, model, optimizer, epoch, train_writer)
train_writer.add_scalar('mean EPE', train_EPE, epoch)
train_EPE_list.append(train_EPE)
# evaluate on validation set
with torch.no_grad():
EPE = validate(val_loader, model, epoch, output_writers)
test_EPE_list.append(EPE)
test_writer.add_scalar('mean EPE', EPE, epoch)
scheduler.step()
if best_EPE < 0:
best_EPE = EPE
is_best = EPE < best_EPE
best_EPE = min(EPE, best_EPE)
save_checkpoint({
'epoch': epoch + 1,
'arch': "photonets",
#'state_dict': model.module.state_dict(),
'best_EPE': best_EPE,
'div_flow': args.div_flow
}, is_best, save_path)
#if args.save_model:
######################
# needs to be updated
######################
#print(train_losses)
#with open("train_losses_one.txt", "wb") as fp: # Pickling
# pickle.dump(train_losses, fp)
#print(test_losses)
#with open("test_losses_one.txt", "wb") as fp: # Pickling
# pickle.dump(test_losses, fp)
print(train_EPE_list)
print(test_EPE_list)
train_EPE_ls = np.asarray(train_EPE_list)
test_EPE_ls = np.asarray(test_EPE_list)
np.save(os.path.join(save_path, "train_EPE_ls"), train_EPE_ls)
np.save(os.path.join(save_path, "test_EPE_ls"), test_EPE_ls)
fig = plt.figure()
plt.plot(epoch_list, train_EPE_list, 'r', label="train")
plt.plot(epoch_list, test_EPE_list, 'b', label="test")
plt.title("End Point Error with Epoch")
plt.legend()
plt.show()
torch.save(model.state_dict(), os.path.join(save_path, "datset1_10epochs_small_minibatch"))
def main():
global args, save_path
args = parser.parse_args()
if args.output_value == 'both':
output_string = "raw output and RGB visualization"
elif args.output_value == 'raw':
output_string = "raw output"
elif args.output_value == 'vis':
output_string = "RGB visualization"
print("=> will save " + output_string)
data_dir = Path(args.data)
print("=> fetching img pairs in '{}'".format(args.data))
if args.output is None:
save_path = data_dir / 'flow'
else:
save_path = Path(args.output)
print('=> will save everything to {}'.format(save_path))
save_path.makedirs_p()
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_pairs = []
for ext in args.img_exts:
test_files = data_dir.files('*0.{}'.format(ext))
for file in test_files:
img_pair = file.parent / (file.namebase[:-1] + '1.{}'.format(ext))
if img_pair.isfile():
img_pairs.append([file, img_pair])
print('{} samples found'.format(len(img_pairs)))
# create model
network_data = torch.load(Path(args.pretrained),
map_location=torch.device('cpu'))
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).to(device)
model.eval()
cudnn.benchmark = True
print("load model success!")
if 'div_flow' in network_data.keys():
args.div_flow = network_data['div_flow']
for (img1_file, img2_file) in tqdm(img_pairs):
img1 = input_transform(imread(img1_file))
img2 = input_transform(imread(img2_file))
input_var = torch.cat([img1, img2]).unsqueeze(0)
print("read image success!")
if args.bidirectional:
# feed inverted pair along with normal pair
inverted_input_var = torch.cat([img2, img1]).unsqueeze(0)
input_var = torch.cat([input_var, inverted_input_var])
input_var = input_var.to(device)
print("prepare to compute flow!")
# compute output
output = model(input_var)
print("flow computing success!")
if args.upsampling is not None:
output = F.interpolate(output,
size=img1.size()[-2:],
mode=args.upsampling,
align_corners=False)
for suffix, flow_output in zip(['flow', 'inv_flow'], output):
filename = save_path / '{}{}'.format(img1_file.namebase[:-1],
suffix)
if args.output_value in ['vis', 'both']:
rgb_flow = flow2rgb(args.div_flow * flow_output,
max_value=args.max_flow)
to_save = (rgb_flow * 255).astype(np.uint8).transpose(1, 2, 0)
imwrite(filename + '.png', to_save)
if args.output_value in ['raw', 'both']:
# Make the flow map a HxWx3 array as in .flo files
to_save = (args.div_flow *
flow_output).cpu().numpy().transpose(1, 2, 0)
np.save(filename + '.npy', to_save)
def main():
global args, best_EPE
args = parser.parse_args()
#checkpoints and model_args
save_path = '{},{},{}epochs{},b{},lr{}'.format(
args.arch, args.solver, args.epochs,
',epochSize' + str(args.epoch_size) if args.epoch_size > 0 else '',
args.batch_size, args.lr)
if not args.no_date:
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
save_path = os.path.join(timestamp, save_path)
save_path = os.path.join(args.dataset, save_path)
print('=> will save everything to {}'.format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
#tensorboardX
train_writer = SummaryWriter(
os.path.join(save_path, 'train')
) #'KITTI_occ/05-29-11:36/flownets,adam,300epochs,epochSize1000,b8,lr0.0001'
test_writer = SummaryWriter(os.path.join(save_path, 'test')) #
output_writers = []
for i in range(3):
output_writers.append(
SummaryWriter(os.path.join(save_path, 'test', str(i))))
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0], std=[args.div_flow, args.div_flow])
])
if 'KITTI' in args.dataset:
args.sparse = True
if args.sparse:
co_transform = flow_transforms.Compose([
flow_transforms.RandomCrop((320, 448)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
else:
co_transform = flow_transforms.Compose([
flow_transforms.RandomTranslate(10),
flow_transforms.RandomRotate(10, 5),
flow_transforms.RandomCrop((320, 448)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
print("=> fetching img pairs in '{}'".format(args.data))
train_set, test_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
target_transform=target_transform,
co_transform=co_transform,
split=args.split_file if args.split_file else args.split_value)
print('{} samples found, {} train samples and {} test samples '.format(
len(test_set) + len(train_set), len(train_set), len(test_set)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=True)
val_loader = torch.utils.data.DataLoader(test_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=False)
# create model 这里的套路可以借鉴下
model = FlowNetS()
if args.pretrained:
network_data = torch.load(args.pretrained)
model.load_state_dict(network_data).to(device)
args.arch = network_data['arch'] #flownets_bn
print("=> using pre-trained model '{}'".format(args.arch))
else:
#network_data = None
#model.load_state_dict(network_data).to(device)
model.init_weights()
print("=> creating model '{}'".format(args.arch))
#这里直接把网络结构都载入了。。。
#model = models.__dict__[args.arch](network_data).to(device)
model = torch.nn.DataParallel(model).to(device) # for multi-GPU
cudnn.benchmark = True
#model settings
#train settings
assert (args.solver in ['adam', 'sgd'])
print('=> setting {} solver'.format(args.solver))
param_groups = [{
'params': model.module.bias_parameters(),
'weight_decay': args.bias_decay
}, {
'params': model.module.weight_parameters(),
'weight_decay': args.weight_decay
}]
if args.solver == 'adam':
optimizer = torch.optim.Adam(param_groups,
args.lr,
betas=(args.momentum, args.beta))
elif args.solver == 'sgd':
optimizer = torch.optim.SGD(param_groups,
args.lr,
momentum=args.momentum)
#evaluate settings
if args.evaluate:
best_EPE = validate(val_loader, model, 0, output_writers)
return
#Decays the learning rate of each parameter group by gamma once the number of epoch reaches
# one of the milestones. Notice that such decay can happen simultaneously with other changes
# to the learning rate from outside this scheduler. When last_epoch=-1, sets initial lr as lr.
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.milestones, gamma=0.5)
# main cycle for train and test
for epoch in range(args.start_epoch, args.epochs):
scheduler.step()
#1. train for one epoch
train_loss, train_EPE = train(train_loader, model, optimizer, epoch,
train_writer)
train_writer.add_scalar('mean EPE', train_EPE, epoch)
#2. evaluate on validation(test)set
with torch.no_grad(): #len(val_loader) == (total *(1- 0.8))/batch_size
EPE = validate(val_loader, model, epoch, output_writers)
test_writer.add_scalar('mean EPE', EPE, epoch)
#3. record the best EPE
if best_EPE < 0:
best_EPE = EPE
is_best = EPE < best_EPE
best_EPE = min(EPE, best_EPE)
#4. save checkpoint
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.module.state_dict(),
'best_EPE': best_EPE,
'div_flow': args.div_flow
}, is_best, save_path)
##############################################################
parser = argparse.ArgumentParser(description='Conversion')
parser.add_argument('--mp4_fn', default='', type=str, help='input path')
parser.add_argument('--gpu_id', default=1, type=int, help='which gpu')
args = parser.parse_args()
model_dir = '/home/fluongo/code/usc_project/FlowNetPytorch/flownets_EPE1.951.pth.tar'
mp4_fn = args.mp4_fn
gpu_id = args.gpu_id
torch.cuda.set_device(gpu_id)
if __name__ == '__main__':
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
# create model
network_data = torch.load(model_dir, map_location='cpu')
print("=> using pre-trained model '{}'".format(network_data['arch']))
model = models.__dict__[network_data['arch']](network_data).cuda()
model.eval()
cudnn.benchmark = True
# if 'div_flow' in network_data.keys():
# args.div_flow = network_data['div_flow']
mp4_ims = mp4_load(mp4_fn)
def main():
global args, best_EPE, save_path
args = parser.parse_args()
save_path = '{},{},{}epochs{},b{},lr{}'.format(
args.arch, args.solver, args.epochs,
',epochSize' + str(args.epoch_size) if args.epoch_size > 0 else '',
args.batch_size, args.lr)
if not args.no_date:
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
save_path = os.path.join(timestamp, save_path)
save_path = os.path.join(args.dataset, save_path)
print('=> will save everything to {}'.format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
train_writer = SummaryWriter(os.path.join(save_path, 'train'))
test_writer = SummaryWriter(os.path.join(save_path, 'test'))
output_writers = []
for i in range(3):
output_writers.append(
SummaryWriter(os.path.join(save_path, 'test', str(i))))
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0], std=[args.div_flow, args.div_flow])
])
if 'KITTI' in args.dataset:
args.sparse = True
if args.sparse:
co_transform = flow_transforms.Compose([
flow_transforms.RandomCrop((320, 448)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
else:
co_transform = flow_transforms.Compose([
flow_transforms.RandomTranslate(10),
flow_transforms.RandomRotate(10, 5),
flow_transforms.RandomCrop((320, 448)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
print("=> fetching img pairs in '{}'".format(args.data))
train_set, test_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
target_transform=target_transform,
co_transform=co_transform,
split=args.split_file if args.split_file else args.split_value)
print('{} samples found, {} train samples and {} test samples '.format(
len(test_set) + len(train_set), len(train_set), len(test_set)))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=True)
val_loader = torch.utils.data.DataLoader(test_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=False)
# create model
if args.pretrained:
network_data = torch.load(args.pretrained)
args.arch = network_data['arch']
print("=> using pre-trained model '{}'".format(args.arch))
else:
network_data = None
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch](network_data).cuda()
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
assert (args.solver in ['adam', 'sgd'])
print('=> setting {} solver'.format(args.solver))
param_groups = [{
'params': model.module.bias_parameters(),
'weight_decay': args.bias_decay
}, {
'params': model.module.weight_parameters(),
'weight_decay': args.weight_decay
}]
if args.solver == 'adam':
optimizer = torch.optim.Adam(param_groups,
args.lr,
betas=(args.momentum, args.beta))
elif args.solver == 'sgd':
optimizer = torch.optim.SGD(param_groups,
args.lr,
momentum=args.momentum)
if args.evaluate:
best_EPE = validate(val_loader, model, 0, output_writers)
return
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.milestones, gamma=0.5)
for epoch in range(args.start_epoch, args.epochs):
scheduler.step()
# train for one epoch
train_loss, train_EPE = train(train_loader, model, optimizer, epoch,
train_writer)
train_writer.add_scalar('mean EPE', train_EPE, epoch)
# evaluate on validation set
EPE = validate(val_loader, model, epoch, output_writers)
test_writer.add_scalar('mean EPE', EPE, epoch)
if best_EPE < 0:
best_EPE = EPE
is_best = EPE < best_EPE
best_EPE = min(EPE, best_EPE)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.module.state_dict(),
'best_EPE': best_EPE,
'div_flow': args.div_flow
}, is_best)
def main():
global args, best_EPE
args = parser.parse_args()
save_path = "{},{},{}epochs{},b{},lr{}".format(
args.arch,
args.solver,
args.epochs,
",epochSize" + str(args.epoch_size) if args.epoch_size > 0 else "",
args.batch_size,
args.lr,
)
if not args.no_date:
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
save_path = os.path.join(timestamp, save_path)
save_path = os.path.join(args.dataset, save_path)
print("=> will save everything to {}".format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
train_writer = SummaryWriter(os.path.join(save_path, "train"))
test_writer = SummaryWriter(os.path.join(save_path, "test"))
output_writers = []
for i in range(3):
output_writers.append(
SummaryWriter(os.path.join(save_path, "test", str(i))))
# Data loading code
if args.data_loader == "torch":
print("Using default data loader \n")
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.45, 0.432, 0.411], std=[1, 1, 1]),
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0],
std=[args.div_flow, args.div_flow]),
])
test_transform = transforms.Compose([
transforms.Resize((122, 162)),
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.45, 0.432, 0.411], std=[1, 1, 1]),
])
if "KITTI" in args.dataset:
args.sparse = True
if args.sparse:
co_transform = flow_transforms.Compose([
flow_transforms.RandomCrop((122, 162)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip(),
])
else:
co_transform = flow_transforms.Compose([
flow_transforms.RandomTranslate(10),
flow_transforms.RandomRotate(10, 5),
flow_transforms.RandomCrop((122, 162)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip(),
])
print("=> fetching img pairs in '{}'".format(args.data))
train_set, test_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
test_transform=test_transform,
target_transform=target_transform,
co_transform=co_transform,
split=args.split_file if args.split_file else args.split_value,
)
print("{} samples found, {} train samples and {} test samples ".format(
len(test_set) + len(train_set), len(train_set), len(test_set)))
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=True,
)
val_loader = torch.utils.data.DataLoader(
test_set,
batch_size=args.batch_size,
num_workers=args.workers,
pin_memory=True,
shuffle=False,
)
if args.data_loader == "dali":
print("Using NVIDIA DALI \n")
(
(image0_train_names, image0_val_names),
(image1_train_names, image1_val_names),
(flow_train_names, flow_val_names),
) = make_dali_dataset(
args.data,
split=args.split_file if args.split_file else args.split_value)
print("{} samples found, {} train samples and {} test samples ".format(
len(image0_val_names) + len(image0_train_names),
len(image0_train_names),
len(image0_val_names),
))
global train_length
global val_length
train_length = len(image0_train_names)
val_length = len(image0_val_names)
def create_image_pipeline(
batch_size,
num_threads,
device_id,
image0_list,
image1_list,
flow_list,
valBool,
):
pipeline = Pipeline(batch_size, num_threads, device_id, seed=2)
with pipeline:
if valBool:
shuffleBool = False
else:
shuffleBool = True
""" READ FILES """
image0, _ = fn.readers.file(
file_root=args.data,
files=image0_list,
random_shuffle=shuffleBool,
name="Reader",
seed=1,
)
image1, _ = fn.readers.file(
file_root=args.data,
files=image1_list,
random_shuffle=shuffleBool,
seed=1,
)
flo = fn.readers.numpy(
file_root=args.data,
files=flow_list,
random_shuffle=shuffleBool,
seed=1,
)
""" DECODE AND RESHAPE """
image0 = fn.decoders.image(image0, device="cpu")
image0 = fn.reshape(image0, layout="HWC")
image1 = fn.decoders.image(image1, device="cpu")
image1 = fn.reshape(image1, layout="HWC")
images = fn.cat(image0, image1, axis=2)
flo = fn.reshape(flo, layout="HWC")
if valBool:
images = fn.resize(images, resize_x=162, resize_y=122)
else:
""" CO-TRANSFORM """
# random translate
# angle_rng = fn.random.uniform(range=(-90, 90))
# images = fn.rotate(images, angle=angle_rng, fill_value=0)
# flo = fn.rotate(flo, angle=angle_rng, fill_value=0)
images = fn.random_resized_crop(
images,
size=[122, 162], # 122, 162
random_aspect_ratio=[1.3, 1.4],
random_area=[0.8, 0.9],
seed=1,
)
flo = fn.random_resized_crop(
flo,
size=[122, 162],
random_aspect_ratio=[1.3, 1.4],
random_area=[0.8, 0.9],
seed=1,
)
# coin1 = fn.random.coin_flip(dtype=types.DALIDataType.BOOL, seed=10)
# coin1_n = coin1 ^ True
# coin2 = fn.random.coin_flip(dtype=types.DALIDataType.BOOL, seed=20)
# coin2_n = coin2 ^ True
# images = (
# fn.flip(images, horizontal=1, vertical=1) * coin1 * coin2
# + fn.flip(images, horizontal=1) * coin1 * coin2_n
# + fn.flip(images, vertical=1) * coin1_n * coin2
# + images * coin1_n * coin2_n
# )
# flo = (
# fn.flip(flo, horizontal=1, vertical=1) * coin1 * coin2
# + fn.flip(flo, horizontal=1) * coin1 * coin2_n
# + fn.flip(flo, vertical=1) * coin1_n * coin2
# + flo * coin1_n * coin2_n
# )
# _flo = flo
# flo_0 = fn.slice(_flo, axis_names="C", start=0, shape=1)
# flo_1 = fn.slice(_flo, axis_names="C", start=1, shape=1)
# flo_0 = flo_0 * coin1 * -1 + flo_0 * coin1_n
# flo_1 = flo_1 * coin2 * -1 + flo_1 * coin2_n
# # flo = noflip + vertical flip + horizontal flip + both_flip
# # A horizontal flip is around the vertical axis (switch left and right)
# # So for a vertical flip coin1 is activated and needs to give +1, coin2 is activated needs to give -1
# # for a horizontal flip coin1 is activated and needs to be -1, coin2_n needs +1
# # no flip coin coin1_n +1, coin2_n +1
# flo = fn.cat(flo_0, flo_1, axis_name="C")
""" NORMALIZE """
images = fn.crop_mirror_normalize(
images,
mean=[0, 0, 0, 0, 0, 0],
std=[255, 255, 255, 255, 255, 255])
images = fn.crop_mirror_normalize(
images,
mean=[0.45, 0.432, 0.411, 0.45, 0.432, 0.411],
std=[1, 1, 1, 1, 1, 1],
)
flo = fn.crop_mirror_normalize(
flo, mean=[0, 0], std=[args.div_flow, args.div_flow])
pipeline.set_outputs(images, flo)
return pipeline
class DALILoader:
def __init__(
self,
batch_size,
image0_names,
image1_names,
flow_names,
valBool,
num_threads,
device_id,
):
self.pipeline = create_image_pipeline(
batch_size,
num_threads,
device_id,
image0_names,
image1_names,
flow_names,
valBool,
)
self.pipeline.build()
self.epoch_size = self.pipeline.epoch_size(
"Reader") / batch_size
output_names = ["images", "flow"]
if valBool:
self.dali_iterator = pytorch.DALIGenericIterator(
self.pipeline,
output_names,
reader_name="Reader",
last_batch_policy=pytorch.LastBatchPolicy.PARTIAL,
auto_reset=True,
)
else:
self.dali_iterator = pytorch.DALIGenericIterator(
self.pipeline,
output_names,
reader_name="Reader",
last_batch_policy=pytorch.LastBatchPolicy.PARTIAL,
auto_reset=True,
)
def __len__(self):
return int(self.epoch_size)
def __iter__(self):
return self.dali_iterator.__iter__()
def reset(self):
return self.dali_iterator.reset()
train_loader = DALILoader(
batch_size=args.batch_size,
num_threads=args.workers,
device_id=0,
image0_names=image0_train_names,
image1_names=image1_train_names,
flow_names=flow_train_names,
valBool=False,
)
val_loader = DALILoader(
batch_size=args.batch_size,
num_threads=args.workers,
device_id=0,
image0_names=image0_val_names,
image1_names=image1_val_names,
flow_names=flow_val_names,
valBool=True,
)
# create model
if args.pretrained:
network_data = torch.load(args.pretrained)
args.arch = network_data["arch"]
print("=> using pre-trained model '{}'".format(args.arch))
else:
network_data = None
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch](network_data).to(device)
assert args.solver in ["adam", "sgd"]
print("=> setting {} solver".format(args.solver))
param_groups = [
{
"params": model.bias_parameters(),
"weight_decay": args.bias_decay
},
{
"params": model.weight_parameters(),
"weight_decay": args.weight_decay
},
]
if device.type == "cuda":
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
if args.solver == "adam":
optimizer = torch.optim.Adam(param_groups,
args.lr,
betas=(args.momentum, args.beta))
elif args.solver == "sgd":
optimizer = torch.optim.SGD(param_groups,
args.lr,
momentum=args.momentum)
if args.evaluate:
best_EPE = validate(val_loader, model, 0, output_writers)
return
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.milestones, gamma=0.5)
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
# # --- quant
# model.train()
# model.qconfig = torch.quantization.get_default_qat_qconfig('qnnpack') # torch.quantization.default_qconfig
# # model = torch.quantization.fuse_modules(model, [['Conv2d', 'bn', 'relu']])
# torch.backends.quantized.engine = 'qnnpack'
# model = torch.quantization.prepare_qat(model)
# # --- quant
# my_sample = next(itertools.islice(train_loader, 10, None))
# print(my_sample[1][0])
# print("Maximum value is ", torch.max(my_sample[0][0]))
# print("Minimum value is ", torch.min(my_sample[0][0]))
train_loss, train_EPE = train(train_loader, model, optimizer, epoch,
train_writer)
train_writer.add_scalar("mean EPE", train_EPE, epoch)
scheduler.step()
# evaluate on validation set
with torch.no_grad():
EPE = validate(val_loader, model, epoch, output_writers)
test_writer.add_scalar("mean EPE", EPE, epoch)
if best_EPE < 0:
best_EPE = EPE
is_best = EPE < best_EPE
best_EPE = min(EPE, best_EPE)
# if is_best:
# kernels = model.module.conv3_1[0].weight.data
# kernels = kernels.cpu()
# kernels = kernels - kernels.min()
# kernels = kernels / kernels.max()
# img = make_grid(kernels)
# plt.imshow(img.permute(1, 2, 0))
# plt.show()
save_checkpoint(
{
"epoch": epoch + 1,
"arch": args.arch,
"state_dict": model.module.state_dict(),
"best_EPE": best_EPE,
"div_flow": args.div_flow,
},
is_best,
save_path,
model,
dummy_input,
)
def test(args, model, img_paths, save_path, idx):
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_file = img_paths[idx]
load_path = img_file
imgId = os.path.basename(img_file)[:-4]
# may get 4 channel (alpha channel) for some format
img_ = imread(load_path)[:, :, :3]
H, W, _ = img_.shape
H_, W_ = int(np.ceil(H / 16.) * 16), int(np.ceil(W / 16.) * 16)
# get spixel id
n_spixl_h = int(np.floor(H_ / args.downsize))
n_spixl_w = int(np.floor(W_ / args.downsize))
spix_values = np.int32(
np.arange(0, n_spixl_w * n_spixl_h).reshape((n_spixl_h, n_spixl_w)))
spix_idx_tensor_ = shift9pos(spix_values)
spix_idx_tensor = np.repeat(np.repeat(spix_idx_tensor_,
args.downsize,
axis=1),
args.downsize,
axis=2)
spixeIds = torch.from_numpy(np.tile(spix_idx_tensor, (1, 1, 1, 1))).type(
torch.float).cuda()
n_spixel = int(n_spixl_h * n_spixl_w)
img = cv2.resize(img_, (W_, H_), interpolation=cv2.INTER_CUBIC)
img1 = input_transform(img)
ori_img = input_transform(img_)
# compute output
tic = time.time()
output = model(img1.cuda().unsqueeze(0))
toc = time.time() - tic
# assign the spixel map
curr_spixl_map = update_spixl_map(spixeIds, output)
ori_sz_spixel_map = F.interpolate(curr_spixl_map.type(torch.float),
size=(H_, W_),
mode='nearest').type(torch.int)
mean_values = torch.tensor([0.411, 0.432, 0.45],
dtype=img1.cuda().unsqueeze(0).dtype).view(
3, 1, 1)
spixel_viz, spixel_label_map = get_spixel_image(
(ori_img + mean_values).clamp(0, 1),
ori_sz_spixel_map.squeeze(),
n_spixels=n_spixel,
b_enforce_connect=True)
# ************************ Save all result********************************************
# save img, uncomment it if needed
# if not os.path.isdir(os.path.join(save_path, 'img')):
# os.makedirs(os.path.join(save_path, 'img'))
# spixl_save_name = os.path.join(save_path, 'img', imgId + '.jpg')
# img_save = (ori_img + mean_values).clamp(0, 1)
# imsave(spixl_save_name, img_save.detach().cpu().numpy().transpose(1, 2, 0))
# save spixel viz
if not os.path.isdir(os.path.join(save_path, 'spixel_viz')):
os.makedirs(os.path.join(save_path, 'spixel_viz'))
spixl_save_name = os.path.join(save_path, 'spixel_viz',
imgId + '_sPixel.png')
imsave(spixl_save_name, spixel_viz.transpose(1, 2, 0))
# save the unique maps as csv, uncomment it if needed
# if not os.path.isdir(os.path.join(save_path, 'map_csv')):
# os.makedirs(os.path.join(save_path, 'map_csv'))
# output_path = os.path.join(save_path, 'map_csv', imgId + '.csv')
# # plus 1 to make it consistent with the toolkit format
# np.savetxt(output_path, (spixel_label_map + 1).astype(int), fmt='%i',delimiter=",")
if idx % 10 == 0:
print("processing %d" % idx)
return toc
def main():
global args, best_EPE
args = parser.parse_args()
if not args.data:
f = open('train_src/data_loc.json', 'r')
content = f.read()
f.close()
data_loc = json.loads(content)
args.data = data_loc[args.dataset]
if not args.savpath:
save_path = '{},{},{}epochs{},b{},lr{}'.format(
args.arch,
args.solver,
args.epochs,
',epochSize'+str(args.epoch_size) if args.epoch_size > 0 else '',
args.batch_size,
args.lr)
if not args.no_date:
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
save_path = os.path.join(timestamp,save_path)
else:
save_path = args.savpath
save_path = os.path.join(args.dataset,save_path)
print('=> will save everything to {}'.format(save_path))
if not os.path.exists(save_path):
os.makedirs(save_path)
# save training args
save_training_args(save_path, args)
train_writer = SummaryWriter(os.path.join(save_path,'train'))
test_writer = SummaryWriter(os.path.join(save_path,'test'))
output_writers = []
for i in range(3):
output_writers.append(SummaryWriter(os.path.join(save_path,'test',str(i))))
# Data loading code
if args.grayscale:
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Grayscale(num_output_channels=3),
transforms.Normalize(mean=[0,0,0], std=[255,255,255]),
transforms.Normalize(mean=[0.431,0.431,0.431], std=[1,1,1]) # 0.431=(0.45+0.432+0.411)/3
# transforms.Normalize(mean=[0.5,0.5,0.5], std=[1,1,1])
])
else:
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0,0], std=[255,255,255]),
transforms.Normalize(mean=[0.45,0.432,0.411], std=[1,1,1])
])
target_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0,0],std=[args.div_flow,args.div_flow])
])
if 'KITTI' in args.dataset:
args.sparse = True
if args.sparse:
co_transform = flow_transforms.Compose([
flow_transforms.RandomCrop((320,448)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
else:
co_transform = flow_transforms.Compose([
flow_transforms.RandomTranslate(10),
flow_transforms.RandomRotate(10,5),
flow_transforms.RandomCrop((320,448)),
flow_transforms.RandomVerticalFlip(),
flow_transforms.RandomHorizontalFlip()
])
print("=> fetching img pairs in '{}'".format(args.data))
train_set, test_set = datasets.__dict__[args.dataset](
args.data,
transform=input_transform,
target_transform=target_transform,
co_transform=co_transform,
split=args.split_file if args.split_file else args.split_value
)
print('{} samp-les found, {} train samples and {} test samples '.format(len(test_set)+len(train_set),
len(train_set),
len(test_set)))
if not args.evaluate:
train_loader = torch.utils.data.DataLoader(
train_set, batch_size=args.batch_size,
num_workers=args.workers, pin_memory=True, shuffle=True)
val_loader = torch.utils.data.DataLoader(
test_set, batch_size=args.batch_size,
num_workers=args.workers, pin_memory=True, shuffle=False)
# create model
if args.pretrained:
network_data = torch.load(args.pretrained)
# args.arch = network_data['arch']
print("=> using pre-trained model '{}'".format(args.arch))
else:
network_data = None
print("=> creating model '{}'".format(args.arch))
# if (args.qw and args.qa and args.cut_ratio) is not None:
# model = models.__dict__[args.arch](data=network_data, bitW=args.qw, bitA=args.qa, cut_ratio=args.cut_ratio).cuda()
# elif (args.qw and args.qa) is not None:
# model = models.__dict__[args.arch](data=network_data, bitW=args.qw, bitA=args.qa).cuda()
# else:
# model = models.__dict__[args.arch](data=network_data).cuda()
model = models.__dict__[args.arch](data=network_data, args=args).to(device)
# model = torch.nn.DataParallel(model).cuda()
# cudnn.benchmark = True
assert(args.solver in ['adam', 'sgd', 'adamw'])
print('=> setting {} solver'.format(args.solver))
param_groups = [{'params': model.bias_parameters(), 'weight_decay': args.bias_decay},
{'params': model.weight_parameters(), 'weight_decay': args.weight_decay}]
if device.type == "cuda":
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
if args.solver == 'adam':
optimizer = torch.optim.Adam(param_groups, args.lr,
betas=(args.momentum, args.beta))
elif args.solver == 'sgd':
optimizer = torch.optim.SGD(param_groups, args.lr,
momentum=args.momentum)
elif args.solver == 'adamw':
optimizer = torch.optim.AdamW(param_groups, args.lr,
betas=(args.momentum, args.beta))
if args.print_model:
exportpars(model, save_path, args)
exportsummary(model, save_path, args)
if args.savpath == 'test':
return
if args.evaluate:
best_EPE = validate(val_loader, model, 0, output_writers)
return
if args.demo:
demo(val_loader, model, 0, output_writers)
return
if args.demovideo:
demovideo(val_loader, model, 0, output_writers)
return
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=args.gamma)
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
train_loss, train_EPE = train(train_loader, model, optimizer, epoch, train_writer)
train_writer.add_scalar('mean EPE', train_EPE, epoch)
scheduler.step()
# evaluate on validation set
with torch.no_grad():
EPE = validate(val_loader, model, epoch, output_writers)
test_writer.add_scalar('mean EPE', EPE, epoch)
if best_EPE < 0:
best_EPE = EPE
is_best = EPE < best_EPE
best_EPE = min(EPE, best_EPE)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.module.state_dict(),
'best_EPE': best_EPE,
'div_flow': args.div_flow
}, is_best, save_path)
def test(model, img_paths, save_path, spixeIds, idx, scale):
# Data loading code
input_transform = transforms.Compose([
flow_transforms.ArrayToTensor(),
transforms.Normalize(mean=[0, 0, 0], std=[255, 255, 255]),
transforms.Normalize(mean=[0.411, 0.432, 0.45], std=[1, 1, 1])
])
img_file = img_paths[idx]
load_path = str(img_file)
imgId = os.path.basename(img_file)[:-4]
# origin size 481*321 or 321*481
img_ = imread(load_path)
H_, W_, _ = img_.shape
# choose the right spixelIndx
spixl_map_idx_tensor = spixeIds[0]
img = cv2.resize(img_, (int(1024 * scale), int(512 * scale)),
interpolation=cv2.INTER_CUBIC)
# if H_ == 321 and W_==481:
# spixl_map_idx_tensor = spixeIds[0]
# img = cv2.resize(img_, (int(480 * scale), int(320 * scale)), interpolation=cv2.INTER_CUBIC)
# elif H_ == 481 and W_ == 321:
# spixl_map_idx_tensor = spixeIds[1]
# img = cv2.resize(img_, (int(320 * scale), int(480 * scale)), interpolation=cv2.INTER_CUBIC)
# else:
# print('The image size is wrong!')
# return
img1 = input_transform(img)
ori_img = input_transform(img_)
mean_values = torch.tensor([0.411, 0.432, 0.45],
dtype=img1.cuda().unsqueeze(0).dtype).view(
3, 1, 1)
# compute output
tic = time.time()
output = model(img1.cuda().unsqueeze(0))
# assign the spixel map and resize to the original size
curr_spixl_map = update_spixl_map(spixl_map_idx_tensor, output)
ori_sz_spixel_map = F.interpolate(curr_spixl_map.type(torch.float),
size=(H_, W_),
mode='nearest').type(torch.int)
spix_index_np = ori_sz_spixel_map.squeeze().detach().cpu().numpy(
).transpose(0, 1)
spix_index_np = spix_index_np.astype(np.int64)
segment_size = (spix_index_np.shape[0] *
spix_index_np.shape[1]) / (int(600 * scale * scale) * 1.0)
min_size = int(0.06 * segment_size)
max_size = int(3 * segment_size)
spixel_label_map = enforce_connectivity(spix_index_np[None, :, :],
min_size, max_size)[0]
torch.cuda.synchronize()
toc = time.time() - tic
n_spixel = len(np.unique(spixel_label_map))
given_img_np = (ori_img + mean_values).clamp(
0, 1).detach().cpu().numpy().transpose(1, 2, 0)
spixel_bd_image = mark_boundaries(given_img_np / np.max(given_img_np),
spixel_label_map.astype(int),
color=(0, 1, 1))
spixel_viz = spixel_bd_image.astype(np.float32).transpose(2, 0, 1)
# ************************ Save all result********************************************
# save img, uncomment it if needed
# if not os.path.isdir(os.path.join(save_path, 'img')):
# os.makedirs(os.path.join(save_path, 'img'))
# spixl_save_name = os.path.join(save_path, 'img', imgId + '.jpg')
# img_save = (ori_img + mean_values).clamp(0, 1)
# imsave(spixl_save_name, img_save.detach().cpu().numpy().transpose(1, 2, 0))
# save spixel viz
if not os.path.isdir(os.path.join(save_path, 'spixel_viz')):
os.makedirs(os.path.join(save_path, 'spixel_viz'))
spixl_save_name = os.path.join(save_path, 'spixel_viz',
imgId + '_sPixel.png')
imsave(spixl_save_name, spixel_viz.transpose(1, 2, 0))
# save the unique maps as csv for eval
# if not os.path.isdir(os.path.join(save_path, 'map_csv')):
# os.makedirs(os.path.join(save_path, 'map_csv'))
# output_path = os.path.join(save_path, 'map_csv', imgId + '.csv')
# # plus 1 to make it consistent with the toolkit format
# np.savetxt(output_path, (spixel_label_map + 1).astype(int), fmt='%i', delimiter=",")
if idx % 10 == 0:
print("processing %d" % idx)
return toc, n_spixel