printer.py cleanup

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_l_up = kp[y][l_hand] < kp[y][l_shoulder]
-    is_r_up = kp[y][r_hand] < kp[y][r_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:

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

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, *,

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:
+        self._front_loop(iterator, axes, number, colors, annotations, dic_out) 
-                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
 
         # Draw the bird figure
         number['num'] = 97
-        for idx in iterator:
+        self._bird_loop(iterator, axes, colors, number)
-            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
         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):
 

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)
+        time.sleep(interval)