Add files via upload

openpifpaf/encoder/__init__.py

@@ -0,0 +1,9 @@
+"""Convert a set of keypoint coordinates into target fields."""
+
+from .encoder import Encoder
+from .factory import cli, factory
+from .paf import Paf
+from .pif import Pif
+from .crm import Crm
+from .skeleton import Skeleton
+from .visualizer import Visualizer

openpifpaf/encoder/annrescaler.py

@@ -0,0 +1,60 @@
+import numpy as np
+
+
+class AnnRescaler(object):
+    def __init__(self, input_output_scale, n_keypoints):
+        self.input_output_scale = input_output_scale
+        self.n_keypoints = n_keypoints
+
+    def __call__(self, anns, width_height_original):
+        keypoint_sets = self.anns_to_keypoint_sets(anns)
+        keypoint_sets[:, :, :2] /= self.input_output_scale
+
+        # background mask
+        bg_mask = self.anns_to_bg_mask(width_height_original, anns)
+        bg_mask = bg_mask[::self.input_output_scale, ::self.input_output_scale]
+
+        # valid area
+        valid_area = None
+        if anns and 'valid_area' in anns[0]:
+            valid_area = anns[0]['valid_area']
+            valid_area = (
+                valid_area[0] / self.input_output_scale,
+                valid_area[1] / self.input_output_scale,
+                valid_area[2] / self.input_output_scale,
+                valid_area[3] / self.input_output_scale,
+            )
+
+        return keypoint_sets, bg_mask, valid_area
+
+    def anns_to_keypoint_sets(self, anns):
+        """Ignore annotations of crowds."""
+        keypoint_sets = [ann['keypoints'] for ann in anns if not ann['iscrowd']]
+        if not keypoint_sets:
+            return np.zeros((0, self.n_keypoints, 3))
+
+        return np.stack(keypoint_sets)
+
+    @staticmethod
+    def anns_to_bg_mask(width_height, anns, include_annotated=True):
+        """Create background mask taking crowded annotations into account."""
+        mask = np.ones(width_height[::-1], dtype=np.bool)
+        for ann in anns:
+            if include_annotated and \
+            not ann['iscrowd'] and \
+            'keypoints' in ann and \
+            np.any(ann['keypoints'][:, 2] > 0):
+                continue
+
+            if 'mask' not in ann:
+                bb = ann['bbox'].copy()
+                bb[2:] += bb[:2]  # convert width and height to x2 and y2
+                bb[0] = np.clip(bb[0], 0, mask.shape[1] - 1)
+                bb[1] = np.clip(bb[1], 0, mask.shape[0] - 1)
+                bb[2] = np.clip(bb[2], 0, mask.shape[1] - 1)
+                bb[3] = np.clip(bb[3], 0, mask.shape[0] - 1)
+                bb = bb.astype(np.int)
+                mask[bb[1]:bb[3] + 1, bb[0]:bb[2] + 1] = 0
+                continue
+            mask[ann['mask']] = 0
+        return mask

openpifpaf/encoder/crm.py

@@ -0,0 +1,33 @@
+import logging
+import re
+
+import numpy as np
+import scipy.ndimage
+import torch
+
+from .annrescaler import AnnRescaler
+from .encoder import Encoder
+from ..utils import create_sink, mask_valid_area
+
+# Not in use because the pre-processing for crm is in datasets.py
+class Crm(Encoder):
+
+    def __init__(self, head_name, stride, *, n_keypoints=None, **kwargs):
+        return
+
+    @staticmethod
+    def match(head_name):
+        return head_name in (
+            'crm',
+        )
+
+    @classmethod
+    def cli(cls, parser):
+        return
+
+    @classmethod
+    def apply_args(cls, args):
+        return
+
+    def __call__(self, anns, width_height_original):
+        return

openpifpaf/encoder/encoder.py

@@ -0,0 +1,19 @@
+from abc import ABCMeta, abstractmethod, abstractstaticmethod
+
+
+class Encoder(metaclass=ABCMeta):
+    @abstractstaticmethod
+    def match(head_name):  # pylint: disable=unused-argument
+        return False
+
+    @classmethod
+    def cli(cls, parser):
+        """Add decoder specific command line arguments to the parser."""
+
+    @classmethod
+    def apply_args(cls, args):
+        """Read command line arguments args to set class properties."""
+
+    @abstractmethod
+    def __call__(self, fields):
+        pass

openpifpaf/encoder/factory.py

@@ -0,0 +1,40 @@
+import logging
+
+from .encoder import Encoder
+from .skeleton import Skeleton
+
+
+def cli(parser):
+    for encoder in Encoder.__subclasses__():
+        encoder.cli(parser)
+
+
+def factory(args, strides):
+    for encoder in Encoder.__subclasses__():
+        encoder.apply_args(args)
+
+    headnames = args.headnets
+
+    #print(strides) 8
+
+    encoders = [factory_head(head_name, stride) for head_name, stride in zip(headnames, strides)]
+    if headnames[-1] == 'skeleton' and len(headnames) == len(strides) + 1:
+        encoders.append(Skeleton())
+
+    return encoders
+
+
+def factory_head(head_name, stride):
+
+    #print("encoder ", Encoder.__subclasses__())
+    #encoder  [<class 'openpifpaf.encoder.paf.Paf'>, <class 'openpifpaf.encoder.pif.Pif'>, <class 'openpifpaf.encoder.crm.Crm'>]
+
+    for encoder in Encoder.__subclasses__():
+        logging.debug('checking whether encoder %s matches %s',
+                      encoder.__name__, head_name)
+        if not encoder.match(head_name):
+            continue
+        logging.info('selected encoder %s for %s', encoder.__name__, head_name)
+        return encoder(head_name, stride)
+
+    raise Exception('unknown head to create an encoder: {}'.format(head_name))

openpifpaf/encoder/paf.py

@@ -0,0 +1,207 @@
+import logging
+import numpy as np
+import scipy
+import torch
+
+from ..data import COCO_PERSON_SKELETON, DENSER_COCO_PERSON_SKELETON, KINEMATIC_TREE_SKELETON
+from .annrescaler import AnnRescaler
+from .encoder import Encoder
+from ..utils import create_sink, mask_valid_area
+
+
+class Paf(Encoder):
+    default_min_size = 3
+    default_fixed_size = False
+    default_aspect_ratio = 0.0
+
+    def __init__(self, head_name, stride, *, skeleton=None, n_keypoints=17, **kwargs):
+        self.log = logging.getLogger(self.__class__.__name__)
+        self.log.debug('unused arguments in %s: %s', head_name, kwargs)
+
+        if skeleton is None:
+            if head_name in ('paf', 'paf19', 'pafs', 'wpaf', 'pafb'):
+                skeleton = COCO_PERSON_SKELETON
+            elif head_name in ('paf16',):
+                skeleton = KINEMATIC_TREE_SKELETON
+            elif head_name in ('paf44',):
+                skeleton = DENSER_COCO_PERSON_SKELETON
+            else:
+                raise Exception('unknown skeleton type of head')
+
+        self.stride = stride
+        self.n_keypoints = n_keypoints
+        self.skeleton = skeleton
+
+        self.min_size = self.default_min_size
+        self.fixed_size = self.default_fixed_size
+        self.aspect_ratio = self.default_aspect_ratio
+
+        if self.fixed_size:
+            assert self.aspect_ratio == 0.0
+
+    @staticmethod
+    def match(head_name):
+        return head_name in (
+            'paf',
+            'paf19',
+            'paf16',
+            'paf44',
+            'pafs',
+            'wpaf',
+            'pafb',
+        )
+
+    @classmethod
+    def cli(cls, parser):
+        group = parser.add_argument_group('paf encoder')
+        group.add_argument('--paf-min-size', default=cls.default_min_size, type=int,
+                           help='min side length of the PAF field')
+        group.add_argument('--paf-fixed-size', default=cls.default_fixed_size, action='store_true',
+                           help='fixed paf size')
+        group.add_argument('--paf-aspect-ratio', default=cls.default_aspect_ratio, type=float,
+                           help='paf width relative to its length')
+
+    @classmethod
+    def apply_args(cls, args):
+        cls.default_min_size = args.paf_min_size
+        cls.default_fixed_size = args.paf_fixed_size
+        cls.default_aspect_ratio = args.paf_aspect_ratio
+
+    def __call__(self, anns, width_height_original):
+        rescaler = AnnRescaler(self.stride, self.n_keypoints)
+        keypoint_sets, bg_mask, valid_area = rescaler(anns, width_height_original)
+        self.log.debug('valid area: %s, paf min size = %d', valid_area, self.min_size)
+
+        f = PafGenerator(self.min_size, self.skeleton,
+                         fixed_size=self.fixed_size, aspect_ratio=self.aspect_ratio)
+        f.init_fields(bg_mask)
+        f.fill(keypoint_sets)
+        return f.fields(valid_area)
+
+
+class PafGenerator(object):
+    def __init__(self, min_size, skeleton, *,
+                 v_threshold=0, padding=10, fixed_size=False, aspect_ratio=0.0):
+        self.min_size = min_size
+        self.skeleton = skeleton
+        self.v_threshold = v_threshold
+        self.padding = padding
+        self.fixed_size = fixed_size
+        self.aspect_ratio = aspect_ratio
+
+        self.intensities = None
+        self.fields_reg1 = None
+        self.fields_reg2 = None
+        self.fields_scale = None
+        self.fields_reg_l = None
+
+    def init_fields(self, bg_mask):
+        n_fields = len(self.skeleton)
+        field_w = bg_mask.shape[1] + 2 * self.padding
+        field_h = bg_mask.shape[0] + 2 * self.padding
+        self.intensities = np.zeros((n_fields + 1, field_h, field_w), dtype=np.float32)
+        self.fields_reg1 = np.zeros((n_fields, 2, field_h, field_w), dtype=np.float32)
+        self.fields_reg2 = np.zeros((n_fields, 2, field_h, field_w), dtype=np.float32)
+        self.fields_scale = np.zeros((n_fields, field_h, field_w), dtype=np.float32)
+        self.fields_reg_l = np.full((n_fields, field_h, field_w), np.inf, dtype=np.float32)
+
+        # set background
+        self.intensities[-1] = 1.0
+        self.intensities[-1, self.padding:-self.padding, self.padding:-self.padding] = bg_mask
+        self.intensities[-1] = scipy.ndimage.binary_erosion(self.intensities[-1],
+                                                            iterations=int(self.min_size / 2.0) + 1,
+                                                            border_value=1)
+
+    def fill(self, keypoint_sets):
+        for keypoints in keypoint_sets:
+            self.fill_keypoints(keypoints)
+
+    def fill_keypoints(self, keypoints):
+        visible = keypoints[:, 2] > 0
+        if not np.any(visible):
+            return
+
+        area = (
+            (np.max(keypoints[visible, 0]) - np.min(keypoints[visible, 0])) *
+            (np.max(keypoints[visible, 1]) - np.min(keypoints[visible, 1]))
+        )
+        scale = np.sqrt(area)
+
+        for i, (joint1i, joint2i) in enumerate(self.skeleton):
+            joint1 = keypoints[joint1i - 1]
+            joint2 = keypoints[joint2i - 1]
+            if joint1[2] <= self.v_threshold or joint2[2] <= self.v_threshold:
+                continue
+
+            self.fill_association(i, joint1, joint2, scale)
+
+    def fill_association(self, i, joint1, joint2, scale):
+        # offset between joints
+        offset = joint2[:2] - joint1[:2]
+        offset_d = np.linalg.norm(offset)
+
+        # dynamically create s
+        s = max(self.min_size, int(offset_d * self.aspect_ratio))
+        sink = create_sink(s)
+        s_offset = (s - 1.0) / 2.0
+
+        # pixel coordinates of top-left joint pixel
+        joint1ij = np.round(joint1[:2] - s_offset)
+        joint2ij = np.round(joint2[:2] - s_offset)
+        offsetij = joint2ij - joint1ij
+
+        # set fields
+        num = max(2, int(np.ceil(offset_d)))
+        fmargin = min(0.4, (s_offset + 1) / (offset_d + np.spacing(1)))
+        # fmargin = 0.0
+        frange = np.linspace(fmargin, 1.0-fmargin, num=num)
+        if self.fixed_size:
+            frange = [0.5]
+        for f in frange:
+            fij = np.round(joint1ij + f * offsetij) + self.padding
+            fminx, fminy = int(fij[0]), int(fij[1])
+            fmaxx, fmaxy = fminx + s, fminy + s
+            if fminx < 0 or fmaxx > self.intensities.shape[2] or \
+               fminy < 0 or fmaxy > self.intensities.shape[1]:
+                continue
+            fxy = (fij - self.padding) + s_offset
+
+            # precise floating point offset of sinks
+            joint1_offset = (joint1[:2] - fxy).reshape(2, 1, 1)
+            joint2_offset = (joint2[:2] - fxy).reshape(2, 1, 1)
+
+            # update intensity
+            self.intensities[i, fminy:fmaxy, fminx:fmaxx] = 1.0
+
+            # update background
+            self.intensities[-1, fminy:fmaxy, fminx:fmaxx] = 0.0
+
+            # update regressions
+            sink1 = sink + joint1_offset
+            sink2 = sink + joint2_offset
+            sink_l = np.minimum(np.linalg.norm(sink1, axis=0),
+                                np.linalg.norm(sink2, axis=0))
+            mask = sink_l < self.fields_reg_l[i, fminy:fmaxy, fminx:fmaxx]
+            self.fields_reg1[i, :, fminy:fmaxy, fminx:fmaxx][:, mask] = \
+                sink1[:, mask]
+            self.fields_reg2[i, :, fminy:fmaxy, fminx:fmaxx][:, mask] = \
+                sink2[:, mask]
+            self.fields_reg_l[i, fminy:fmaxy, fminx:fmaxx][mask] = sink_l[mask]
+
+            # update scale
+            self.fields_scale[i, fminy:fmaxy, fminx:fmaxx][mask] = scale
+
+    def fields(self, valid_area):
+        intensities = self.intensities[:, self.padding:-self.padding, self.padding:-self.padding]
+        fields_reg1 = self.fields_reg1[:, :, self.padding:-self.padding, self.padding:-self.padding]
+        fields_reg2 = self.fields_reg2[:, :, self.padding:-self.padding, self.padding:-self.padding]
+        fields_scale = self.fields_scale[:, self.padding:-self.padding, self.padding:-self.padding]
+
+        mask_valid_area(intensities, valid_area)
+
+        return (
+            torch.from_numpy(intensities),
+            torch.from_numpy(fields_reg1),
+            torch.from_numpy(fields_reg2),
+            torch.from_numpy(fields_scale),
+        )