diff --git a/monoloco/activity.py b/monoloco/activity.py index 3c8f46e..f13b695 100644 --- a/monoloco/activity.py +++ b/monoloco/activity.py @@ -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: diff --git a/monoloco/run.py b/monoloco/run.py index 7f66d7a..b22497a 100644 --- a/monoloco/run.py +++ b/monoloco/run.py @@ -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') diff --git a/monoloco/visuals/pifpaf_show.py b/monoloco/visuals/pifpaf_show.py index dbf65ac..5fef1d2 100644 --- a/monoloco/visuals/pifpaf_show.py +++ b/monoloco/visuals/pifpaf_show.py @@ -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, *, diff --git a/monoloco/visuals/printer.py b/monoloco/visuals/printer.py index 555f3c6..34fce8a 100644 --- a/monoloco/visuals/printer.py +++ b/monoloco/visuals/printer.py @@ -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): diff --git a/monoloco/visuals/webcam.py b/monoloco/visuals/webcam.py index d8778ff..8edab16 100644 --- a/monoloco/visuals/webcam.py +++ b/monoloco/visuals/webcam.py @@ -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) \ No newline at end of file + time.sleep(interval)