[WIP][DO NOT MERGE] Matrixnet

gluoncv/data/transforms/presets/__init__.py

@@ -6,3 +6,4 @@
 from . import imagenet
 from . import simple_pose
 from . import segmentation
+from . import matrix_net

gluoncv/data/transforms/presets/matrix_net.py

@@ -0,0 +1,196 @@
+"""Transforms described in https://arxiv.org/abs/1904.07850 and https://arxiv.org/abs/2001.03194."""
+# pylint: disable=too-many-function-args
+from __future__ import absolute_import
+import numpy as np
+import mxnet as mx
+from .. import bbox as tbbox
+from .. import image as timage
+from .. import experimental
+from ....utils.filesystem import try_import_cv2
+
+__all__ = ['MatrixNetDefaultTrainTransform', 'MatrixNetDefaultValTransform',
+ 'get_post_transform']
+
+class MatrixNetDefaultTrainTransform(object):
+ """Default MatrixNet training transform which includes tons of image augmentations.
+
+ Parameters
+ ----------
+ width : int
+ Image width.
+ height : int
+ Image height.
+ num_class : int
+ Number of categories
+ layers_range : list of list of number(list of number)
+ Represents the same meaning as that of MatrixNet
+ scale_factor : int, default is 4
+ The downsampling scale factor between input image and output heatmap
+ mean : array-like of size 3
+ Mean pixel values to be subtracted from image tensor. Default is [0.485, 0.456, 0.406].
+ std : array-like of size 3
+ Standard deviation to be divided from image. Default is [0.229, 0.224, 0.225].
+ """
+ def __init__(self, width, height, num_class, layers_range, scale_factor=4, mean=(0.485, 0.456, 0.406),
+ std=(0.229, 0.224, 0.225), **kwargs):
+ self._kwargs = kwargs
+ self._width = width
+ self._height = height
+ self._num_class = num_class
+ self._layers_range = layers_range
+ self._scale_factor = scale_factor
+ self._mean = np.array(mean, dtype=np.float32).reshape(1, 1, 3)
+ self._std = np.array(std, dtype=np.float32).reshape(1, 1, 3)
+ self._data_rng = np.random.RandomState(123)
+ self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571],
+ dtype=np.float32)
+ self._eig_vec = np.array([
+ [-0.58752847, -0.69563484, 0.41340352],
+ [-0.5832747, 0.00994535, -0.81221408],
+ [-0.56089297, 0.71832671, 0.41158938]
+ ], dtype=np.float32)
+
+ from ....model_zoo.matrix_net.target_generator import MatrixNetTargetGenerator
+ self._target_generator = MatrixNetTargetGenerator(
+ num_class, width, height, self._layers_range)
+
+ def __call__(self, src, label):
+ """Apply transform to training image/label."""
+ # random color jittering
+ img = src
+ bbox = label
+
+ # random horizontal flip
+ h, w, _ = img.shape
+ img, flips = timage.random_flip(img, px=0.5)
+ bbox = tbbox.flip(bbox, (w, h), flip_x=flips[0])
+
+ cv2 = try_import_cv2()
+ input_h, input_w = self._height, self._width
+ s = max(h, w) * 1.0
+ c = np.array([w / 2., h / 2.], dtype=np.float32)
+ sf = 0.4
+ w_border = _get_border(128, img.shape[1])
+ h_border = _get_border(128, img.shape[0])
+ c[0] = np.random.randint(low=w_border, high=img.shape[1] - w_border)
+ c[1] = np.random.randint(low=h_border, high=img.shape[0] - h_border)
+ s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)
+ trans_input = tbbox.get_affine_transform(c, s, 0, [input_w, input_h])
+ inp = cv2.warpAffine(img.asnumpy(), trans_input, (input_w, input_h), flags=cv2.INTER_LINEAR)
+
+ trans_output = tbbox.get_affine_transform(c, s, 0, [input_w, input_h])
+ for i in range(bbox.shape[0]):
+ bbox[i, :2] = tbbox.affine_transform(bbox[i, :2], trans_output)
+ bbox[i, 2:4] = tbbox.affine_transform(bbox[i, 2:4], trans_output)
+ bbox[:, :2] = np.clip(bbox[:, :2], 0, input_w - 1)
+ bbox[:, 2:4] = np.clip(bbox[:, 2:4], 0, input_h - 1)
+ img = inp
+
+ # to tensor
+ img = img.astype(np.float32) / 255.
+ experimental.image.np_random_color_distort(img, data_rng=self._data_rng)
+ img = (img - self._mean) / self._std
+ img = img.transpose(2, 0, 1).astype(np.float32)
+ img = mx.nd.array(img)
+
+ # generate training target so cpu workers can help reduce the workload on gpu
+ gt_bboxes = bbox[:, :4]
+ gt_ids = bbox[:, 4:5]
+ heatmaps, wh_targets, wh_masks, center_regs, center_reg_masks = self._target_generator(
+ gt_bboxes, gt_ids)
+ results = []
+ results.append(img)
+ for heatmap in heatmaps:
+ results.append(heatmap)
+ for wh_target in wh_targets:
+ results.append(wh_target)
+ for wh_mask in wh_masks:
+ results.append(wh_mask)
+ for center_reg in center_regs:
+ results.append(center_reg)
+ for center_reg_mask in center_reg_masks:
+ results.append(center_reg_mask)
+ return tuple(results)
+
+
+class MatrixNetDefaultValTransform(object):
+ """Default MatrixNet validation transform.
+
+ Parameters
+ ----------
+ width : int
+ Image width.
+ height : int
+ Image height.
+ mean : array-like of size 3
+ Mean pixel values to be subtracted from image tensor. Default is [0.485, 0.456, 0.406].
+ std : array-like of size 3
+ Standard deviation to be divided from image. Default is [0.229, 0.224, 0.225].
+
+ """
+ def __init__(self, width, height, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+ self._width = width
+ self._height = height
+ self._mean = np.array(mean, dtype=np.float32).reshape(1, 1, 3)
+ self._std = np.array(std, dtype=np.float32).reshape(1, 1, 3)
+
+ def __call__(self, src, label):
+ """Apply transform to validation image/label."""
+ # resize
+ img, bbox = src.asnumpy(), label
+ cv2 = try_import_cv2()
+ input_h, input_w = self._height, self._width
+ h, w, _ = src.shape
+ s = max(h, w) * 1.0
+ c = np.array([w / 2., h / 2.], dtype=np.float32)
+ trans_input = tbbox.get_affine_transform(c, s, 0, [input_w, input_h])
+ inp = cv2.warpAffine(img, trans_input, (input_w, input_h), flags=cv2.INTER_LINEAR)
+ output_w = input_w
+ output_h = input_h
+ trans_output = tbbox.get_affine_transform(c, s, 0, [output_w, output_h])
+ for i in range(bbox.shape[0]):
+ bbox[i, :2] = tbbox.affine_transform(bbox[i, :2], trans_output)
+ bbox[i, 2:4] = tbbox.affine_transform(bbox[i, 2:4], trans_output)
+ bbox[:, :2] = np.clip(bbox[:, :2], 0, output_w - 1)
+ bbox[:, 2:4] = np.clip(bbox[:, 2:4], 0, output_h - 1)
+ img = inp
+
+ # to tensor
+ img = img.astype(np.float32) / 255.
+ img = (img - self._mean) / self._std
+ img = img.transpose(2, 0, 1).astype(np.float32)
+ img = mx.nd.array(img)
+ return img, bbox.astype(img.dtype)
+
+def get_post_transform(orig_w, orig_h, out_w, out_h):
+ """Get the post prediction affine transforms. This will be used to adjust the prediction results
+ according to original coco image resolutions.
+
+ Parameters
+ ----------
+ orig_w : int
+ Original width of the image.
+ orig_h : int
+ Original height of the image.
+ out_w : int
+ Width of the output image after prediction.
+ out_h : int
+ Height of the output image after prediction.
+
+ Returns
+ -------
+ numpy.ndarray
+ Affine transform matrix 3x2.
+
+ """
+ s = max(orig_w, orig_h) * 1.0
+ c = np.array([orig_w / 2., orig_h / 2.], dtype=np.float32)
+ trans_output = tbbox.get_affine_transform(c, s, 0, [out_w, out_h], inv=True)
+ return trans_output
+
+def _get_border(border, size):
+ """Get the border size of the image"""
+ i = 1
+ while size - border // i <= border // i:
+ i *= 2
+ return border // i

gluoncv/model_zoo/matrix_net/__init__.py

@@ -0,0 +1,5 @@
+"""MatrixNet"""
+# pylint: disable=wildcard-import
+from __future__ import absolute_import
+
+from .matrix_net import *