printer.py cleanup
This commit is contained in:
parent
e64ab138b3
commit
3458cc58e9
@ -10,7 +10,9 @@ import torch
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from .network.process import laplace_sampling
|
||||
from .visuals.pifpaf_show import KeypointPainter, image_canvas, get_pifpaf_outputs, draw_orientation, social_distance_colors
|
||||
from .visuals.pifpaf_show import (
|
||||
KeypointPainter, image_canvas, get_pifpaf_outputs, draw_orientation, social_distance_colors
|
||||
)
|
||||
|
||||
|
||||
def social_interactions(idx, centers, angles, dds, stds=None, social_distance=False,
|
||||
@ -85,23 +87,23 @@ def is_raising_hand(kp):
|
||||
|
||||
head_width = kp[x][l_ear]- kp[x][r_ear]
|
||||
head_top = (kp[y][nose] - head_width)
|
||||
|
||||
|
||||
l_forearm = [kp[x][l_hand] - kp[x][l_elbow], kp[y][l_hand] - kp[y][l_elbow]]
|
||||
l_arm = [kp[x][l_shoulder] - kp[x][l_elbow], kp[y][l_shoulder] - kp[y][l_elbow]]
|
||||
l_arm = [kp[x][l_shoulder] - kp[x][l_elbow], kp[y][l_shoulder] - kp[y][l_elbow]]
|
||||
|
||||
r_forearm = [kp[x][r_hand] - kp[x][r_elbow], kp[y][r_hand] - kp[y][r_elbow]]
|
||||
r_arm = [kp[x][r_shoulder] - kp[x][r_elbow], kp[y][r_shoulder] - kp[y][r_elbow]]
|
||||
r_arm = [kp[x][r_shoulder] - kp[x][r_elbow], kp[y][r_shoulder] - kp[y][r_elbow]]
|
||||
|
||||
l_angle = (90/np.pi) * np.arccos(np.dot(l_forearm/np.linalg.norm(l_forearm), l_arm/np.linalg.norm(l_arm)))
|
||||
r_angle = (90/np.pi) * np.arccos(np.dot(r_forearm/np.linalg.norm(r_forearm), r_arm/np.linalg.norm(r_arm)))
|
||||
|
||||
is_l_up = kp[y][l_hand] < kp[y][l_shoulder]
|
||||
is_r_up = kp[y][r_hand] < kp[y][r_shoulder]
|
||||
|
||||
is_l_up = kp[y][l_hand] < kp[y][l_shoulder]
|
||||
is_r_up = kp[y][r_hand] < kp[y][r_shoulder]
|
||||
|
||||
l_too_close = kp[x][l_hand] <= kp[x][l_shoulder] and kp[y][l_hand]>=head_top
|
||||
r_too_close = kp[x][r_hand] >= kp[x][r_shoulder] and kp[y][r_hand]>=head_top
|
||||
|
||||
is_left_risen = is_l_up and l_angle >= 30 and not l_too_close
|
||||
is_left_risen = is_l_up and l_angle >= 30 and not l_too_close
|
||||
is_right_risen = is_r_up and r_angle >= 30 and not r_too_close
|
||||
|
||||
if is_left_risen and is_right_risen:
|
||||
|
||||
@ -48,7 +48,8 @@ def cli():
|
||||
visualizer.cli(parser)
|
||||
|
||||
# Monoloco
|
||||
predict_parser.add_argument('--activities', nargs='+', help='Choose activities to show: social_distance, raise_hand')
|
||||
predict_parser.add_argument('--activities', nargs='+',
|
||||
help='Choose activities to show: social_distance, raise_hand')
|
||||
predict_parser.add_argument('--mode', help='keypoints, mono, stereo', default='mono')
|
||||
predict_parser.add_argument('--model', help='path of MonoLoco/MonStereo model to load')
|
||||
predict_parser.add_argument('--net', help='only to select older MonoLoco model, otherwise use --mode')
|
||||
|
||||
@ -100,7 +100,7 @@ class KeypointPainter:
|
||||
width, height = (1,1)
|
||||
if size:
|
||||
width = size[0]
|
||||
height = size[1]
|
||||
height = size[1]
|
||||
|
||||
l_arm_width = np.sqrt(((x[9]-x[7])/width)**2 + ((y[9]-y[7])/height)**2)*100
|
||||
r_arm_width = np.sqrt(((x[10]-x[8])/width)**2 + ((y[10]-y[8])/height)**2)*100
|
||||
@ -109,10 +109,12 @@ class KeypointPainter:
|
||||
for ci, connection in enumerate(np.array(self.skeleton) - 1):
|
||||
c = color
|
||||
linewidth=self.linewidth
|
||||
if ((connection[0] == 5 and connection[1] == 7) or (connection[0] == 7 and connection[1] == 9)) and raise_hand in ['left','both']:
|
||||
if ((connection[0] == 5 and connection[1] == 7)
|
||||
or (connection[0] == 7 and connection[1] == 9)) and raise_hand in ['left','both']:
|
||||
c = 'yellow'
|
||||
linewidth = l_arm_width
|
||||
if ((connection[0] == 6 and connection[1] == 8) or (connection[0] == 8 and connection[1] == 10)) and raise_hand in ['right', 'both']:
|
||||
if ((connection[0] == 6 and connection[1] == 8)
|
||||
or (connection[0] == 8 and connection[1] == 10)) and raise_hand in ['right', 'both']:
|
||||
c = 'yellow'
|
||||
linewidth = r_arm_width
|
||||
if self.color_connections:
|
||||
@ -190,7 +192,8 @@ class KeypointPainter:
|
||||
matplotlib.patches.Rectangle(
|
||||
(x - scale, y - scale), 2 * scale, 2 * scale, fill=False, color=color))
|
||||
|
||||
def keypoints(self, ax, keypoint_sets, *, size=None, scores=None, color=None, colors=None, texts=None, raise_hand='none'):
|
||||
def keypoints(self, ax, keypoint_sets, *,
|
||||
size=None, scores=None, color=None, colors=None, texts=None, raise_hand='none'):
|
||||
if keypoint_sets is None:
|
||||
return
|
||||
|
||||
@ -211,7 +214,7 @@ class KeypointPainter:
|
||||
if isinstance(color, (int, np.integer)):
|
||||
color = matplotlib.cm.get_cmap('tab20')((color % 20 + 0.05) / 20)
|
||||
|
||||
if raise_hand is not 'none':
|
||||
if raise_hand != 'none':
|
||||
# if raise_hand[:][i] is 'both' or raise_hand[:][i] is 'left' or raise_hand[:][i] is 'right':
|
||||
# color = 'green'
|
||||
self._draw_skeleton(ax, x, y, v, size=size, color=color, raise_hand=raise_hand[:][i])
|
||||
@ -228,19 +231,6 @@ class KeypointPainter:
|
||||
|
||||
if texts is not None:
|
||||
self._draw_text(ax, x, y, v, texts[i], color)
|
||||
|
||||
# nose = 0
|
||||
# l_ear = 3
|
||||
# l_shoulder = 5
|
||||
# r_ear = 4
|
||||
# r_shoulder = 6
|
||||
|
||||
# head_width = kps[l_ear][0]- kps[r_ear][0]
|
||||
# head_top = (kps[nose][1] - head_width)
|
||||
|
||||
# ax.plot([kps[l_shoulder][0],kps[l_shoulder][0]], [kps[l_shoulder][1],head_top], linewidth=10, color='red')
|
||||
# ax.plot([kps[r_shoulder][0],kps[r_shoulder][0]], [kps[r_shoulder][1],head_top], linewidth=10, color='red')
|
||||
# ax.plot([kps[l_shoulder][0],kps[r_shoulder][0]], [head_top,head_top], linewidth=10, color='red')
|
||||
|
||||
|
||||
def annotations(self, ax, annotations, *,
|
||||
|
||||
@ -139,7 +139,7 @@ class Printer:
|
||||
|
||||
fig, (ax0, ax1) = plt.subplots(1, 2, sharey=False, gridspec_kw={'width_ratios': [width_ratio, 1]},
|
||||
figsize=(fig_width, fig_height))
|
||||
|
||||
|
||||
ax1.set_aspect(fig_ar_1)
|
||||
fig.set_tight_layout(True)
|
||||
fig.subplots_adjust(left=0.02, right=0.98, bottom=0, top=1, hspace=0, wspace=0.02)
|
||||
@ -195,6 +195,39 @@ class Printer:
|
||||
def social_distance_bird(self, axis, colors):
|
||||
draw_orientation(axis, self.xz_centers, [], self.angles, colors, mode='bird')
|
||||
|
||||
|
||||
def _front_loop(self, iterator, axes, number, colors, annotations, dic_out):
|
||||
for idx in iterator:
|
||||
if any(xx in self.output_types for xx in ['front', 'multi']) and self.zz_pred[idx] > 0:
|
||||
if self.args.activities:
|
||||
if 'social_distance' in self.args.activities:
|
||||
self.social_distance_front(axes[0], colors, annotations, dic_out)
|
||||
elif 'raise_hand' in self.args.activities:
|
||||
self.social_distance_front(axes[0], colors, annotations, dic_out)
|
||||
else:
|
||||
self._draw_front(axes[0],
|
||||
self.dd_pred[idx],
|
||||
idx,
|
||||
number)
|
||||
number['num'] += 1
|
||||
|
||||
|
||||
def _bird_loop(self, iterator, axes, colors, number):
|
||||
for idx in iterator:
|
||||
if any(xx in self.output_types for xx in ['bird', 'multi']) and self.zz_pred[idx] > 0:
|
||||
|
||||
if self.args.activities:
|
||||
if 'social_distance' in self.args.activities:
|
||||
self.social_distance_bird(axes[1], colors)
|
||||
# Draw ground truth and uncertainty
|
||||
self._draw_uncertainty(axes, idx)
|
||||
|
||||
# Draw bird eye view text
|
||||
if number['flag']:
|
||||
self._draw_text_bird(axes, idx, number['num'])
|
||||
number['num'] += 1
|
||||
|
||||
|
||||
def draw(self, figures, axes, image, dic_out=None, annotations=None):
|
||||
|
||||
if self.args.activities:
|
||||
@ -211,37 +244,16 @@ class Printer:
|
||||
number = dict(flag=False, num=97)
|
||||
if any(xx in self.output_types for xx in ['front', 'multi']):
|
||||
number['flag'] = True # add numbers
|
||||
# Remove image if social distance is activated
|
||||
if not self.args.activities or 'social_distance' not in self.args.activities:
|
||||
self.mpl_im0.set_data(image)
|
||||
for idx in iterator:
|
||||
if any(xx in self.output_types for xx in ['front', 'multi']) and self.zz_pred[idx] > 0:
|
||||
if self.args.activities:
|
||||
if 'social_distance' in self.args.activities:
|
||||
self.social_distance_front(axes[0], colors, annotations, dic_out)
|
||||
elif 'raise_hand' in self.args.activities:
|
||||
self.social_distance_front(axes[0], colors, annotations, dic_out)
|
||||
else:
|
||||
self._draw_front(axes[0],
|
||||
self.dd_pred[idx],
|
||||
idx,
|
||||
number)
|
||||
number['num'] += 1
|
||||
|
||||
self._front_loop(iterator, axes, number, colors, annotations, dic_out)
|
||||
|
||||
# Draw the bird figure
|
||||
number['num'] = 97
|
||||
for idx in iterator:
|
||||
if any(xx in self.output_types for xx in ['bird', 'multi']) and self.zz_pred[idx] > 0:
|
||||
self._bird_loop(iterator, axes, colors, number)
|
||||
|
||||
if self.args.activities:
|
||||
if 'social_distance' in self.args.activities:
|
||||
self.social_distance_bird(axes[1], colors)
|
||||
# Draw ground truth and uncertainty
|
||||
self._draw_uncertainty(axes, idx)
|
||||
|
||||
# Draw bird eye view text
|
||||
if number['flag']:
|
||||
self._draw_text_bird(axes, idx, number['num'])
|
||||
number['num'] += 1
|
||||
self._draw_legend(axes)
|
||||
|
||||
# Draw, save or/and show the figures
|
||||
@ -254,7 +266,6 @@ class Printer:
|
||||
if self.plt_close:
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
|
||||
def _draw_front(self, ax, z, idx, number):
|
||||
|
||||
|
||||
@ -7,7 +7,6 @@ Implementation adapted from https://github.com/vita-epfl/openpifpaf/blob/master/
|
||||
"""
|
||||
|
||||
import time
|
||||
import os
|
||||
import logging
|
||||
|
||||
import torch
|
||||
@ -36,13 +35,13 @@ def factory_from_args(args):
|
||||
args.checkpoint = dic_models['keypoints']
|
||||
|
||||
logger.configure(args, LOG) # logger first
|
||||
|
||||
|
||||
if args.output_types is None:
|
||||
args.output_types = ['multi']
|
||||
|
||||
assert 'bird' not in args.output_types
|
||||
if 'json' not in args.output_types:
|
||||
assert len(args.output_types) is 1
|
||||
assert len(args.output_types) == 1
|
||||
else:
|
||||
assert len(args.output_types) < 3
|
||||
|
||||
@ -77,8 +76,8 @@ def factory_from_args(args):
|
||||
|
||||
|
||||
def webcam(args):
|
||||
|
||||
assert args.mode in ('mono')
|
||||
|
||||
assert args.mode in 'mono'
|
||||
|
||||
args, dic_models = factory_from_args(args)
|
||||
|
||||
@ -202,4 +201,4 @@ def mypause(interval):
|
||||
canvas.draw_idle()
|
||||
canvas.start_event_loop(interval)
|
||||
else:
|
||||
time.sleep(interval)
|
||||
time.sleep(interval)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user