lint(): run pylint

AmitMY · AmitMY · commit 994afb17f132 · 2021-02-07T17:23:55.000+02:00
diff --git a/LICENSE.txt b/LICENSE.txt
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2020 Amit Moryossef
+Copyright (c) 2021 Amit Moryossef
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/README.md b/README.md
@@ -8,7 +8,7 @@ The format supports any type of poses, arbitrary number of people, and arbitrary
 
 The main idea is having a `header` with instructions on how many points exists, where, and how to connect them.
 
-The binary spec can be found in [format/spec.md](format/spec.md).
+The binary spec can be found in [lib/specs/v0.1.md](lib/specs/v0.1.md).
 
 ### Python Usage
 ```bash
diff --git a/pose_format/pose.py b/pose_format/pose.py
@@ -11,123 +11,135 @@
 
 
 class Pose:
-    """File IO for '.pose' file format, including the header and body"""
-
-    def __init__(self, header: PoseHeader, body: PoseBody):
-        """
-        :param header: PoseHeader
-        :param body: PoseBody
-        """
-        self.header = header
-        self.body = body
-
-    @staticmethod
-    def read(buffer: bytes, pose_body: PoseBody = NumPyPoseBody):
-        reader = BufferReader(buffer)
-        header = PoseHeader.read(reader)
-        body = pose_body.read(header, reader)
-
-        return Pose(header, body)
-
-    def write(self, buffer: BinaryIO):
-        self.header.write(buffer)
-        self.body.write(self.header.version, buffer)
-
-    def focus(self):
-        """
-        Gets the pose to start at (0,0) and have dimensions as big as needed
-        """
-        mins = ma.min(self.body.data, axis=(0, 1, 2))
-        maxs = ma.max(self.body.data, axis=(0, 1, 2))
-
-        if np.count_nonzero(mins) > 0:  # Only translate if there is a number to translate by
-            self.body.data = ma.subtract(self.body.data, mins)
-
-        dimensions = (maxs - mins).tolist()
-        self.header.dimensions = PoseHeaderDimensions(*dimensions)
-
-    def normalize(self, info: PoseNormalizationInfo, scale_factor: float = 1):
-        """
-        Normalize the point to a fixed distance between two points
-        """
-        mask = self.body.data.mask
-        transposed = self.body.zero_filled().points_perspective()
-
-        p1s = transposed[info.p1]
-        p2s = transposed[info.p2]
-
-        if transposed.shape[1] == 0:
-            p1s = p1s[0]
-            p2s = p2s[0]
-        else:
-            p1s = ma.concatenate(p1s)
-            p2s = ma.concatenate(p2s)
-
-        # Move all points so center is (0,0)
-        center = np.mean((p2s + p1s) / 2, axis=0)
-        self.body.data -= center
-
-        mean_distance = np.mean(distance_batch(p1s, p2s))
-
-        scale = scale_factor / mean_distance  # scale all points to dist/scale
-
-        if round(scale, 5) != 1:
-            self.body.data = ma.multiply(self.body.data, scale)
-
-        self.body.data = ma.array(self.body.data, mask=mask)
-
-        return self
-
-    def frame_dropout(self, dropout_std=0.1):
-        body, selected_indexes = self.body.frame_dropout(dropout_std=dropout_std)
-        return Pose(header=self.header, body=body), selected_indexes
-
-    def get_components(self, components: List[str]):
-        indexes = []
-        new_components = []
-
-        idx = 0
-        for component in self.header.components:
-            if component.name in components:
-                new_components.append(component)
-                indexes += list(range(idx, len(component.points) + idx))
-            idx += len(component.points)
-
-        new_header = PoseHeader(self.header.version, self.header.dimensions, new_components)
-        new_body = self.body.get_points(indexes)
-
-        return Pose(header=new_header, body=new_body)
-
-    def bbox(self):
-        body = self.body.bbox(self.header)
-        header = self.header.bbox()
-        return Pose(header=header, body=body)
-
-    pass_through_methods = {
-        "augment2d",  # Augment 2D points
-        "interpolate",  # Interpolate missing pose points
-        "torch",  # Convert body to torch
-        "tensorflow",  # Convert body to tensorflow
-        "slice_step",  # Step through the data
-    }
-
-    def __getattr__(self, attr):
-        if attr not in Pose.pass_through_methods:
-            raise AttributeError("Attribute '%s' doesn't exist on class Pose" % attr)
-
-        def func(*args, **kwargs):
-            prop = getattr(self.body, attr)
-            body_res = prop(*args, **kwargs)
+  """File IO for '.pose' file format, including the header and body"""
 
-            if isinstance(body_res, PoseBody):
-                header = self.header
-                if hasattr(header, attr):
-                    header_res = getattr(header, attr)(*args, **kwargs)
-                    if isinstance(header_res, PoseHeader):
-                        header = header_res
+  def __init__(self, header: PoseHeader, body: PoseBody):
+    """
+    :param header: PoseHeader
+    :param body: PoseBody
+    """
+    self.header = header
+    self.body = body
 
-                return Pose(header, body_res)
+  @staticmethod
+  def read(buffer: bytes, pose_body: PoseBody = NumPyPoseBody):
+    reader = BufferReader(buffer)
+    header = PoseHeader.read(reader)
+    body = pose_body.read(header, reader)
 
-            return body_res
+    return Pose(header, body)
 
-        return func
+  def write(self, buffer: BinaryIO):
+    self.header.write(buffer)
+    self.body.write(self.header.version, buffer)
+
+  def focus(self):
+    """
+    Gets the pose to start at (0,0) and have dimensions as big as needed
+    """
+    mins = ma.min(self.body.data, axis=(0, 1, 2))
+    maxs = ma.max(self.body.data, axis=(0, 1, 2))
+
+    if np.count_nonzero(mins) > 0:  # Only translate if there is a number to translate by
+      self.body.data = ma.subtract(self.body.data, mins)
+
+    dimensions = (maxs - mins).tolist()
+    self.header.dimensions = PoseHeaderDimensions(*dimensions)
+
+  def normalize(self, info: PoseNormalizationInfo, scale_factor: float = 1):
+    """
+    Normalize the point to a fixed distance between two points
+    """
+    mask = self.body.data.mask
+    transposed = self.body.zero_filled().points_perspective()
+
+    p1s = transposed[info.p1]
+    p2s = transposed[info.p2]
+
+    if transposed.shape[1] == 0:
+      p1s = p1s[0]
+      p2s = p2s[0]
+    else:
+      p1s = ma.concatenate(p1s)
+      p2s = ma.concatenate(p2s)
+
+    # Move all points so center is (0,0)
+    center = np.mean((p2s + p1s) / 2, axis=0)
+    self.body.data -= center
+
+    mean_distance = np.mean(distance_batch(p1s, p2s))
+
+    scale = scale_factor / mean_distance  # scale all points to dist/scale
+
+    if round(scale, 5) != 1:
+      self.body.data = ma.multiply(self.body.data, scale)
+
+    self.body.data = ma.array(self.body.data, mask=mask)
+
+    return self
+
+  def normalize_distribution(self, mu=None, std=None):
+    mu = mu if mu is not None else ma.mean(self.body.data, axis=(0, 1))
+    std = std if std is not None else ma.std(self.body.data, axis=(0, 1))
+
+    self.body.data = (self.body.data - mu) / std
+
+    return mu, std
+
+  def unnormalize_distribution(self, mu, std):
+    self.body.data = (self.body.data * std) + mu
+
+  def frame_dropout(self, dropout_std=0.1):
+    body, selected_indexes = self.body.frame_dropout(dropout_std=dropout_std)
+    return Pose(header=self.header, body=body), selected_indexes
+
+  def get_components(self, components: List[str]):
+    indexes = []
+    new_components = []
+
+    idx = 0
+    for component in self.header.components:
+      if component.name in components:
+        new_components.append(component)
+        indexes += list(range(idx, len(component.points) + idx))
+      idx += len(component.points)
+
+    new_header = PoseHeader(self.header.version, self.header.dimensions, new_components)
+    new_body = self.body.get_points(indexes)
+
+    return Pose(header=new_header, body=new_body)
+
+  def bbox(self):
+    body = self.body.bbox(self.header)
+    header = self.header.bbox()
+    return Pose(header=header, body=body)
+
+  pass_through_methods = {
+    "augment2d",  # Augment 2D points
+    "flip",  # Flip pose on axis
+    "interpolate",  # Interpolate missing pose points
+    "torch",  # Convert body to torch
+    "tensorflow",  # Convert body to tensorflow
+    "slice_step",  # Step through the data
+  }
+
+  def __getattr__(self, attr):
+    if attr not in Pose.pass_through_methods:
+      raise AttributeError("Attribute '%s' doesn't exist on class Pose" % attr)
+
+    def func(*args, **kwargs):
+      prop = getattr(self.body, attr)
+      body_res = prop(*args, **kwargs)
+
+      if isinstance(body_res, PoseBody):
+        header = self.header
+        if hasattr(header, attr):
+          header_res = getattr(header, attr)(*args, **kwargs)
+          if isinstance(header_res, PoseHeader):
+            header = header_res
+
+        return Pose(header, body_res)
+
+      return body_res
+
+    return func
diff --git a/pose_format/pose_body.py b/pose_format/pose_body.py
@@ -129,10 +129,14 @@ def augment2d(self, rotation_std=0.2, shear_std=0.2, scale_std=0.2):
         # Based on https://en.wikipedia.org/wiki/Scaling_(geometry)
         if scale_std > 0:
             scale_matrix = np.eye(2)
-            scale_matrix[0][0] += np.random.normal(loc=0, scale=scale_std, size=1)[0]
+            scale_matrix[1][1] += np.random.normal(loc=0, scale=scale_std, size=1)[0]
             matrix = np.dot(matrix, scale_matrix)
 
-        return self.matmul(matrix.astype(dtype=np.float32))
+        # Cast to matrix the correct size
+        dim_matrix = np.eye(self.data.shape[-1])
+        dim_matrix[0:2,0:2] = matrix
+
+        return self.matmul(dim_matrix.astype(dtype=np.float32))
 
     def zero_filled(self) -> __qualname__:
         raise NotImplementedError("'zero_filled' not implemented on '%s'" % self.__class__)
diff --git a/pose_format/pose_visualizer.py b/pose_format/pose_visualizer.py
@@ -1,4 +1,5 @@
 import itertools
+from functools import lru_cache
 from typing import Tuple, Iterator
 
 import cv2
@@ -15,39 +16,42 @@ def __init__(self, pose: Pose):
         self.pose = pose
 
     def _draw_frame(self, frame: ma.MaskedArray, frame_confidence: np.ndarray, img) -> np.ndarray:
-        avg_color = np.mean(img, axis=(0, 1))
-        # print("avg_color", avg_color)
+        background_color = img[0][0] # Estimation of background color for opacity. `mean` is slow
+
+        thickness = round(math.sqrt(img.shape[0] * img.shape[1]) / 300)
+        radius = round(thickness/2)
 
         for person, person_confidence in zip(frame, frame_confidence):
             c = person_confidence.tolist()
             idx = 0
             for component in self.pose.header.components:
                 colors = [np.array(c[::-1]) for c in component.colors]
 
+                @lru_cache(maxsize=None)
                 def _point_color(p_i: int):
                     opacity = c[p_i + idx]
-                    np_color = colors[p_i % len(component.colors)] * opacity + (1 - opacity) * avg_color
+                    np_color = colors[p_i % len(component.colors)] * opacity + (1 - opacity) * background_color
                     return tuple([int(c) for c in np_color])
 
                 # Draw Points
                 for i in range(len(component.points)):
                     if c[i + idx] > 0:
-                        cv2.circle(img=img, center=tuple(person[i + idx]), radius=3,
+                        cv2.circle(img=img, center=tuple(person[i + idx][:2]), radius=radius,
                                    color=_point_color(i), thickness=-1)
 
                 if self.pose.header.is_bbox:
                     point1 = tuple(person[0 + idx].tolist())
                     point2 = tuple(person[1 + idx].tolist())
                     color = tuple(np.mean([_point_color(0), _point_color(1)], axis=0))
 
-                    cv2.rectangle(img=img, pt1=point1, pt2=point2, color=color, thickness=2)
+                    cv2.rectangle(img=img, pt1=point1, pt2=point2, color=color, thickness=thickness)
                 else:
                     int_person = person.astype(np.int32)
                     # Draw Limbs
                     for (p1, p2) in component.limbs:
                         if c[p1 + idx] > 0 and c[p2 + idx] > 0:
-                            point1 = tuple(int_person[p1 + idx].tolist())
-                            point2 = tuple(int_person[p2 + idx].tolist())
+                            point1 = tuple(int_person[p1 + idx].tolist()[:2])
+                            point2 = tuple(int_person[p2 + idx].tolist()[:2])
 
                             length = ((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2) ** 0.5
 
@@ -56,7 +60,7 @@ def _point_color(p_i: int):
                             deg = math.degrees(math.atan2(point1[1] - point2[1], point1[0] - point2[0]))
                             polygon = cv2.ellipse2Poly(
                                 (int((point1[0] + point2[0]) / 2), int((point1[1] + point2[1]) / 2)),
-                                (int(length / 2), 3),
+                                (int(length / 2), thickness),
                                 int(deg),
                                 0, 360, 1)
                             cv2.fillConvexPoly(img=img, points=polygon, color=color)
@@ -67,28 +71,37 @@ def _point_color(p_i: int):
 
     def draw(self, background_color: Tuple[int, int, int] = (255, 255, 255), max_frames: int = None):
         int_data = np.array(np.around(self.pose.body.data.data), dtype="int32")
+        background = np.full((self.pose.header.dimensions.height, self.pose.header.dimensions.width, 3),
+                             fill_value=background_color, dtype="uint8")
         for frame, confidence in itertools.islice(zip(int_data, self.pose.body.confidence), max_frames):
-            background = np.full((self.pose.header.dimensions.height, self.pose.header.dimensions.width, 3),
-                                 fill_value=background_color,
-                                 dtype="uint8")
-            yield self._draw_frame(frame, confidence, img=background)
+            yield self._draw_frame(frame, confidence, img=background.copy())
 
-    def draw_on_video(self, background_video: str, max_frames: int = None, blur=False):
+    def draw_on_video(self, background_video, max_frames: int = None, blur=False):
         int_data = np.array(np.around(self.pose.body.data.data), dtype="int32")
 
         if max_frames is None:
             max_frames = len(int_data)
 
-        cap = cv2.VideoCapture(background_video)
-        for frame, confidence in itertools.islice(zip(int_data, self.pose.body.confidence), max_frames):
-            _, background = cap.read()
+        def get_frames(video_path):
+
+            cap = cv2.VideoCapture(video_path)
+            while True:
+                ret, vf = cap.read()
+                if not ret:
+                    break
+                yield vf
+            cap.release()
+
+        if isinstance(background_video, str):
+            background_video = iter(get_frames(background_video))
+
+        for frame, confidence, background in itertools.islice(zip(int_data, self.pose.body.confidence, background_video), max_frames):
             background = cv2.resize(background, (self.pose.header.dimensions.width, self.pose.header.dimensions.height))
 
             if blur:
                 background = cv2.blur(background, (20, 20))
 
             yield self._draw_frame(frame, confidence, background)
-        cap.release()
 
     def save_frame(self, f_name: str, frame: np.ndarray):
         cv2.imwrite(f_name, frame)
diff --git a/setup.py b/setup.py
