From 286d4d136689b938bb815d2818676cdf34abb23a Mon Sep 17 00:00:00 2001 From: Haziq Razali Date: Mon, 23 Sep 2019 22:38:46 +0800 Subject: [PATCH] Add files via upload --- openpifpaf/encoder/__init__.py | 9 ++ openpifpaf/encoder/annrescaler.py | 60 +++++++++ openpifpaf/encoder/crm.py | 33 +++++ openpifpaf/encoder/encoder.py | 19 +++ openpifpaf/encoder/factory.py | 40 ++++++ openpifpaf/encoder/paf.py | 207 ++++++++++++++++++++++++++++++ openpifpaf/encoder/pif.py | 154 ++++++++++++++++++++++ openpifpaf/encoder/skeleton.py | 23 ++++ openpifpaf/encoder/visualizer.py | 82 ++++++++++++ 9 files changed, 627 insertions(+) create mode 100644 openpifpaf/encoder/__init__.py create mode 100644 openpifpaf/encoder/annrescaler.py create mode 100644 openpifpaf/encoder/crm.py create mode 100644 openpifpaf/encoder/encoder.py create mode 100644 openpifpaf/encoder/factory.py create mode 100644 openpifpaf/encoder/paf.py create mode 100644 openpifpaf/encoder/pif.py create mode 100644 openpifpaf/encoder/skeleton.py create mode 100644 openpifpaf/encoder/visualizer.py diff --git a/openpifpaf/encoder/__init__.py b/openpifpaf/encoder/__init__.py new file mode 100644 index 0000000..eeae8d8 --- /dev/null +++ b/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 diff --git a/openpifpaf/encoder/annrescaler.py b/openpifpaf/encoder/annrescaler.py new file mode 100644 index 0000000..e9287b5 --- /dev/null +++ b/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 diff --git a/openpifpaf/encoder/crm.py b/openpifpaf/encoder/crm.py new file mode 100644 index 0000000..fc82676 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/openpifpaf/encoder/encoder.py b/openpifpaf/encoder/encoder.py new file mode 100644 index 0000000..34cb122 --- /dev/null +++ b/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 diff --git a/openpifpaf/encoder/factory.py b/openpifpaf/encoder/factory.py new file mode 100644 index 0000000..d2a4a3b --- /dev/null +++ b/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 [, , ] + + 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)) diff --git a/openpifpaf/encoder/paf.py b/openpifpaf/encoder/paf.py new file mode 100644 index 0000000..0dbbc93 --- /dev/null +++ b/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), + ) diff --git a/openpifpaf/encoder/pif.py b/openpifpaf/encoder/pif.py new file mode 100644 index 0000000..2f01c84 --- /dev/null +++ b/openpifpaf/encoder/pif.py @@ -0,0 +1,154 @@ +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 + + +class Pif(Encoder): + default_side_length = 4 + + def __init__(self, head_name, stride, *, n_keypoints=None, **kwargs): + self.log = logging.getLogger(self.__class__.__name__) + self.log.debug('unused arguments in %s: %s', head_name, kwargs) + + self.stride = stride + if n_keypoints is None: + m = re.match('pif([0-9]+)$', head_name) + if m is not None: + n_keypoints = int(m.group(1)) + self.log.debug('using %d keypoints for pif', n_keypoints) + else: + n_keypoints = 17 + self.n_keypoints = n_keypoints + self.side_length = self.default_side_length + + @staticmethod + def match(head_name): + return head_name in ( + 'pif', + 'ppif', + 'pifb', + 'pifs', + ) or re.match('pif([0-9]+)$', head_name) is not None + + @classmethod + def cli(cls, parser): + group = parser.add_argument_group('pif encoder') + group.add_argument('--pif-side-length', default=cls.default_side_length, type=int, + help='side length of the PIF field') + + @classmethod + def apply_args(cls, args): + cls.default_side_length = args.pif_side_length + + 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, pif side length = %d', valid_area, self.side_length) + + n_fields = keypoint_sets.shape[1] + f = PifGenerator(self.side_length) + f.init_fields(n_fields, bg_mask) + f.fill(keypoint_sets) + return f.fields(valid_area) + + +class PifGenerator(object): + def __init__(self, side_length, v_threshold=0, padding=10): + self.side_length = side_length + self.v_threshold = v_threshold + self.padding = padding + + self.intensities = None + self.fields_reg = None + self.fields_scale = None + self.fields_reg_l = None + + self.sink = create_sink(side_length) + self.s_offset = (side_length - 1.0) / 2.0 + + self.log = logging.getLogger(self.__class__.__name__) + + def init_fields(self, n_fields, bg_mask): + 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_reg = 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) + + # bg_mask + 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.s_offset) + 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) + self.log.debug('instance scale = %.3f', scale) + + for f, xyv in enumerate(keypoints): + if xyv[2] <= self.v_threshold: + continue + + self.fill_coordinate(f, xyv, scale) + + def fill_coordinate(self, f, xyv, scale): + ij = np.round(xyv[:2] - self.s_offset).astype(np.int) + self.padding + minx, miny = int(ij[0]), int(ij[1]) + maxx, maxy = minx + self.side_length, miny + self.side_length + if minx < 0 or maxx > self.intensities.shape[2] or \ + miny < 0 or maxy > self.intensities.shape[1]: + return + + offset = xyv[:2] - (ij + self.s_offset - self.padding) + offset = offset.reshape(2, 1, 1) + + # update intensity + self.intensities[f, miny:maxy, minx:maxx] = 1.0 + + # allow unknown margin in background + self.intensities[-1, miny:maxy, minx:maxx] = 0.0 + + # update regression + sink_reg = self.sink + offset + sink_l = np.linalg.norm(sink_reg, axis=0) + mask = sink_l < self.fields_reg_l[f, miny:maxy, minx:maxx] + self.fields_reg[f, :, miny:maxy, minx:maxx][:, mask] = \ + sink_reg[:, mask] + self.fields_reg_l[f, miny:maxy, minx:maxx][mask] = sink_l[mask] + + # update scale + self.fields_scale[f, miny:maxy, minx:maxx][mask] = scale + + def fields(self, valid_area): + intensities = self.intensities[:, self.padding:-self.padding, self.padding:-self.padding] + fields_reg = self.fields_reg[:, :, 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_reg), + torch.from_numpy(fields_scale), + ) diff --git a/openpifpaf/encoder/skeleton.py b/openpifpaf/encoder/skeleton.py new file mode 100644 index 0000000..6012d85 --- /dev/null +++ b/openpifpaf/encoder/skeleton.py @@ -0,0 +1,23 @@ +import numpy as np +import torch + +from .annrescaler import AnnRescaler + + +class Skeleton(object): + def __init__(self, *, max_instances=100, n_keypoints=17): + self.max_instances = max_instances + self.n_keypoints = n_keypoints + + def __call__(self, anns, width_height_original, v_th=0): + rescaler = AnnRescaler(1, self.n_keypoints) + keypoint_sets, __, ___ = rescaler(anns, width_height_original) + + out = np.zeros((self.max_instances, self.n_keypoints, 3), dtype=np.float) + for i, keypoints in enumerate(keypoint_sets): + for f, xyv in enumerate(keypoints): + if xyv[2] <= v_th: + continue + out[i, f] = xyv + + return (torch.from_numpy(out),) diff --git a/openpifpaf/encoder/visualizer.py b/openpifpaf/encoder/visualizer.py new file mode 100644 index 0000000..007b402 --- /dev/null +++ b/openpifpaf/encoder/visualizer.py @@ -0,0 +1,82 @@ +import numpy as np + +from ..data import COCO_KEYPOINTS, COCO_PERSON_SKELETON +from .. import show + + +class Visualizer(object): + def __init__(self, headnames, strides): + self.headnames = headnames + self.strides = strides + self.keypoint_painter = show.KeypointPainter(skeleton=COCO_PERSON_SKELETON) + + def single(self, image, targets): + keypoint_sets = None + if 'skeleton' in self.headnames: + i = self.headnames.index('skeleton') + keypoint_sets = targets[i][0] + + with show.canvas() as ax: + ax.imshow(image) + + for target, headname, stride in zip(targets, self.headnames, self.strides): + print(headname, len(target)) + if headname in ('paf', 'paf19', 'pafs', 'wpaf'): + self.paf19(image, target, stride, keypoint_sets) + elif headname in ('pif', 'pif17', 'pifs'): + self.pif17(image, target, stride, keypoint_sets) + + def pif17(self, image, target, stride, keypoint_sets): + resized_image = image[::stride, ::stride] + for f in [1, 2, 15, 16]: + print('intensity field', COCO_KEYPOINTS[f]) + + with show.canvas() as ax: + ax.imshow(resized_image) + bce_mask = np.sum(target[0], axis=0) > 0.0 + ax.imshow(target[0][f] + 0.5 * bce_mask, alpha=0.9, vmin=0.0, vmax=1.0) + + with show.canvas() as ax: + ax.imshow(image) + show.white_screen(ax, alpha=0.5) + self.keypoint_painter.keypoints(ax, keypoint_sets) + show.quiver(ax, target[1][f], xy_scale=stride) + + def paf19(self, image, target, stride, keypoint_sets): + resized_image = image[::stride, ::stride] + for f in [1, 2, 15, 16, 17, 18]: + print('association field', + COCO_KEYPOINTS[COCO_PERSON_SKELETON[f][0] - 1], + COCO_KEYPOINTS[COCO_PERSON_SKELETON[f][1] - 1]) + with show.canvas() as ax: + ax.imshow(resized_image) + bce_mask = np.sum(target[0], axis=0) > 0.0 + ax.imshow(target[0][f] + 0.5 * bce_mask, alpha=0.9, vmin=0.0, vmax=1.0) + + with show.canvas() as ax: + ax.imshow(image) + show.white_screen(ax, alpha=0.5) + self.keypoint_painter.keypoints(ax, keypoint_sets) + show.quiver(ax, target[1][f], xy_scale=stride) + + with show.canvas() as ax: + ax.imshow(image) + show.white_screen(ax, alpha=0.5) + self.keypoint_painter.keypoints(ax, keypoint_sets) + show.quiver(ax, target[2][f], xy_scale=stride) + + def __call__(self, images, targets, meta): + n_heads = len(targets) + n_batch = len(images) + targets = [[t.numpy() for t in heads] for heads in targets] + targets = [ + [[target_field[batch_i] for target_field in targets[head_i]] + for head_i in range(n_heads)] + for batch_i in range(n_batch) + ] + + images = np.moveaxis(np.asarray(images), 1, -1) + images = np.clip((images + 2.0) / 4.0, 0.0, 1.0) + + for i, t in zip(images, targets): + self.single(i, t)