charles/monoloco
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

predict with new pifpaf

Browse Source
This commit is contained in:
Lorenzo 2020-12-22 12:42:32 +01:00
parent c877a16c4b
commit 966b692e4d
4 changed files with 37 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
monstereo/network/net.py
View File

@ -56,7 +56,7 @@ class Loco:
output_size=output_size) output_size=output_size)
else: else:
self.model = MonStereoModel(p_dropout=p_dropout, input_size=input_size, output_size=output_size, self.model = MonStereoModel(p_dropout=p_dropout, input_size=input_size, output_size=output_size,
linear_size=linear_size, device=self.device) linear_size=linear_size, device=self.device)
self.model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage)) self.model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
else: else:

3
monstereo/network/process.py
View File

@ -82,7 +82,7 @@ def factory_for_gt(im_size, name=None, path_gt=None, verbose=True):
dic_gt = None dic_gt = None
x_factor = im_size[0] / 1600 x_factor = im_size[0] / 1600
y_factor = im_size[1] / 900 y_factor = im_size[1] / 900
pixel_factor = (x_factor + y_factor) / 1.75 # 1.7 for MOT pixel_factor = (x_factor + y_factor) / 1.75 # 1.75 for MOT
# pixel_factor = 1 # pixel_factor = 1
if im_size[0] / im_size[1] > 2.5: if im_size[0] / im_size[1] > 2.5:
kk = [[718.3351, 0., 600.3891], [0., 718.3351, 181.5122], [0., 0., 1.]] # Kitti calibration kk = [[718.3351, 0., 600.3891], [0., 718.3351, 181.5122], [0., 0., 1.]] # Kitti calibration
@ -274,7 +274,6 @@ def extract_outputs(outputs, tasks=()):
if outputs.shape[1] == 10: if outputs.shape[1] == 10:
dic_out['aux'] = torch.sigmoid(dic_out['aux']) dic_out['aux'] = torch.sigmoid(dic_out['aux'])
return dic_out return dic_out

58
monstereo/predict.py
View File

@ -105,19 +105,18 @@ def predict(args):
pifpaf_out = [ann.json_data() for ann in pred] pifpaf_out = [ann.json_data() for ann in pred]
if batch_i == 0: if batch_i == 0:
pifpaf_outputs = [keypoint_sets, scores, pifpaf_out] # keypoints_sets and scores for pifpaf printing pifpaf_outputs = pred # to only print left image for stereo
images_outputs = [cpu_image] # List of 1 or 2 elements with pifpaf tensor and monoloco original image
pifpaf_outs = {'left': pifpaf_out} pifpaf_outs = {'left': pifpaf_out}
with open(meta_batch[0]['file_name'], 'rb') as f:
cpu_image = PIL.Image.open(f).convert('RGB')
else: else:
pifpaf_outs['right'] = pifpaf_out pifpaf_outs['right'] = pifpaf_out
# Load the original image # Load the original image
if args.net in ('monoloco_pp', 'monstereo'): if args.net in ('monoloco_pp', 'monstereo'):
with open(meta['file_name'], 'rb') as f:
cpu_image = PIL.Image.open(f).convert('RGB')
im_name = os.path.basename(meta['file_name']) im_name = os.path.basename(meta['file_name'])
im_size = (cpu_image.size()[1], cpu_image.size()[0]) # Original im_size = (cpu_image.size[0], cpu_image.size[1]) # Original
kk, dic_gt = factory_for_gt(im_size, name=im_name, path_gt=args.path_gt) kk, dic_gt = factory_for_gt(im_size, name=im_name, path_gt=args.path_gt)
# Preprocess pifpaf outputs and run monoloco # Preprocess pifpaf outputs and run monoloco
@ -138,53 +137,52 @@ def predict(args):
dic_out = defaultdict(list) dic_out = defaultdict(list)
kk = None kk = None
factory_outputs(args, images_outputs, output_path, pifpaf_outputs, dic_out=dic_out, kk=kk) # TODO Clean
factory_outputs(args, annotation_painter, cpu_image, output_path, pifpaf_outputs, pifpaf_out,
dic_out=dic_out, kk=kk)
print('Image {}\n'.format(cnt) + '-' * 120) print('Image {}\n'.format(cnt) + '-' * 120)
cnt += 1 cnt += 1
def factory_outputs(args, images_outputs, output_path, pifpaf_outputs, dic_out=None, kk=None): def factory_outputs(args, annotation_painter, cpu_image, output_path, pred, pifpaf_out, dic_out=None, kk=None):
"""Output json files or images according to the choice""" """Output json files or images according to the choice"""
# Save json file # Save json file
if args.mode == 'pifpaf': if args.net == 'pifpaf':
with show.image_canvas(cpu_image, image_out_name) as ax: with openpifpaf.show.image_canvas(cpu_image, output_path) as ax:
annotation_painter.annotations(ax, pred) annotation_painter.annotations(ax, pred)
keypoint_sets, scores, pifpaf_out = pifpaf_outputs[:]
# Visualizer # Visualizer
keypoint_painter = KeypointPainter(show_box=False) keypoint_painter = KeypointPainter(show_box=False)
skeleton_painter = KeypointPainter(show_box=False, color_connections=True, markersize=1, linewidth=4) skeleton_painter = KeypointPainter(show_box=False, color_connections=True, markersize=1, linewidth=4)
if 'json' in args.output_types and keypoint_sets.size > 0: if 'json' in args.output_types and len(pred) > 0:
with open(output_path + '.pifpaf.json', 'w') as f: with open(output_path + '.pifpaf.json', 'w') as f:
json.dump(pifpaf_out, f) json.dump(pifpaf_out, f)
if 'keypoints' in args.output_types: # if 'keypoints' in args.output_types:
with image_canvas(images_outputs[0], # with image_canvas(images_outputs[0],
output_path + '.keypoints.png', # output_path + '.keypoints.png',
show=args.show, # show=args.show,
fig_width=args.figure_width, # fig_width=args.figure_width,
dpi_factor=args.dpi_factor) as ax: # dpi_factor=args.dpi_factor) as ax:
keypoint_painter.keypoints(ax, keypoint_sets) # keypoint_painter.keypoints(ax, keypoint_sets)
#
if 'skeleton' in args.output_types: # if 'skeleton' in args.output_types:
with image_canvas(images_outputs[0], # with image_canvas(images_outputs[0],
output_path + '.skeleton.png', # output_path + '.skeleton.png',
show=args.show, # show=args.show,
fig_width=args.figure_width, # fig_width=args.figure_width,
dpi_factor=args.dpi_factor) as ax: # dpi_factor=args.dpi_factor) as ax:
skeleton_painter.keypoints(ax, keypoint_sets, scores=scores) # skeleton_painter.keypoints(ax, keypoint_sets, scores=scores)
else: else:
if any((xx in args.output_types for xx in ['front', 'bird', 'multi'])): if any((xx in args.output_types for xx in ['front', 'bird', 'multi'])):
print(output_path) print(output_path)
if dic_out['boxes']: # Only print in case of detections if dic_out['boxes']: # Only print in case of detections
printer = Printer(images_outputs[1], output_path, kk, args) printer = Printer(cpu_image, output_path, kk, args)
figures, axes = printer.factory_axes() figures, axes = printer.factory_axes()
printer.draw(figures, axes, dic_out, images_outputs[1]) printer.draw(figures, axes, dic_out, cpu_image)
if 'json' in args.output_types: if 'json' in args.output_types:
with open(os.path.join(output_path + '.monoloco.json'), 'w') as ff: with open(os.path.join(output_path + '.monoloco.json'), 'w') as ff:

10
monstereo/visuals/printer.py
View File

@ -86,8 +86,7 @@ class Printer:
self.boxes = dic_ann['boxes'] self.boxes = dic_ann['boxes']
self.boxes_gt = dic_ann['boxes_gt'] self.boxes_gt = dic_ann['boxes_gt']
self.uv_camera = (int(self.im.size[0] / 2), self.im.size[1]) self.uv_camera = (int(self.im.size[0] / 2), self.im.size[1])
if dic_ann['aux']: self.auxs = dic_ann['aux']
self.auxs = dic_ann['aux'] if dic_ann['aux'] else None
def factory_axes(self): def factory_axes(self):
"""Create axes for figures: front bird multi""" """Create axes for figures: front bird multi"""
@ -199,7 +198,12 @@ class Printer:
def _draw_front(self, ax, z, idx, number): def _draw_front(self, ax, z, idx, number):
mode = 'stereo' if self.auxs[idx] > 0.3 else 'mono' if len(self.auxs) == 0:
mode = 'mono'
elif self.auxs[idx] <= 0.3:
mode = 'mono'
else:
mode = 'stereo'
# Bbox # Bbox
w = min(self.width-2, self.boxes[idx][2] - self.boxes[idx][0]) w = min(self.width-2, self.boxes[idx][2] - self.boxes[idx][0])