Huawei_Technology
/
mindspore
Not watched
1
0
Fork 0
Code
Releases 13 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: 098292eefa
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '098292eefa'
${ noResults }
mindspore/model_zoo/official/cv/openpose/src/dataset.py

575 lines
26 kB
Raw Blame History 

  1. # Copyright 2020 Huawei Technologies Co., Ltd

  2. 

  3. # Licensed under the Apache License, Version 2.0 (the "License");

  4. # you may not use this file except in compliance with the License.

  5. # You may obtain a copy of the License at

  6. 

  7. # http://www.apache.org/licenses/LICENSE-2.0

  8. 

  9. # Unless required by applicable law or agreed to in writing, software

  10. # distributed under the License is distributed on an "AS IS" BASIS,

  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  12. # See the License for the specific language governing permissions and

  13. # limitations under the License.

  14. # ============================================================================

  15. import os

  16. import math

  17. import random

  18. import numpy as np

  19. import cv2

  20. from pycocotools.coco import COCO as ReadJson

  21. 

  22. import mindspore.dataset as de

  23. 

  24. from src.config import JointType, params

  25. 

  26. cv2.setNumThreads(0)

  27. 

  28. class txtdataset():

  29.     def __init__(self, train, imgpath, maskpath, insize, mode='train', n_samples=None):

  30.         self.train = train

  31.         self.mode = mode

  32.         self.imgpath = imgpath

  33.         self.maskpath = maskpath

  34.         self.insize = insize

  35.         self.maxtime = 0

  36.         self.catIds = train.getCatIds(catNms=['person'])

  37.         self.imgIds = sorted(train.getImgIds(catIds=self.catIds))

  38.         if self.mode == 'train':

  39.             self.clean_imgIds()

  40.         if self.mode in ['val', 'eval'] and n_samples is not None:

  41.             self.imgIds = random.sample(self.imgIds, n_samples)

  42.         print('{} images: {}'.format(mode, len(self)))

  43. 

  44. 

  45.     def __len__(self):

  46.         return len(self.imgIds)

  47. 

  48.     def clean_imgIds(self):

  49.         print("cleaning imgids")

  50. 

  51.         for img_id in self.imgIds.copy():

  52.             annotations = None

  53.             anno_ids = self.train.getAnnIds(imgIds=[img_id], iscrowd=None)

  54. 

  55.             # annotation for that image

  56.             if anno_ids:

  57.                 annotations_for_img = self.train.loadAnns(anno_ids)

  58. 

  59.                 person_cnt = 0

  60.                 valid_annotations_for_img = []

  61.                 for annotation in annotations_for_img:

  62.                     # if too few keypoints or too small

  63.                     if annotation['num_keypoints'] >= params['min_keypoints'] and \

  64.                             annotation['area'] > params['min_area']:

  65.                         person_cnt += 1

  66.                         valid_annotations_for_img.append(annotation)

  67. 

  68.                 # if person annotation

  69.                 if person_cnt > 0:

  70.                     annotations = valid_annotations_for_img

  71.             if annotations is None:

  72.                 self.imgIds.remove(img_id)

  73. 

  74.     def overlay_paf(self, img, paf):

  75.         hue = ((np.arctan2(paf[1], paf[0]) / np.pi) / -2 + 0.5)

  76.         saturation = np.sqrt(paf[0] ** 2 + paf[1] ** 2)

  77.         saturation[saturation > 1.0] = 1.0

  78.         value = saturation.copy()

  79.         hsv_paf = np.vstack((hue[np.newaxis], saturation[np.newaxis], value[np.newaxis])).transpose(1, 2, 0)

  80.         rgb_paf = cv2.cvtColor((hsv_paf * 255).astype(np.uint8), cv2.COLOR_HSV2BGR)

  81.         img = cv2.addWeighted(img, 0.6, rgb_paf, 0.4, 0)

  82.         return img

  83. 

  84.     def overlay_pafs(self, img, pafs):

  85.         mix_paf = np.zeros((2,) + img.shape[:-1])

  86.         paf_flags = np.zeros(mix_paf.shape) # for constant paf

  87. 

  88.         for paf in pafs.reshape((int(pafs.shape[0]/2), 2,) + pafs.shape[1:]):

  89.             paf_flags = paf != 0

  90.             paf_flags += np.broadcast_to(paf_flags[0] | paf_flags[1], paf.shape)

  91.             mix_paf += paf

  92. 

  93.         mix_paf[paf_flags > 0] /= paf_flags[paf_flags > 0]

  94.         img = self.overlay_paf(img, mix_paf)

  95.         return img

  96. 

  97.     def overlay_heatmap(self, img, heatmap):

  98.         rgb_heatmap = cv2.applyColorMap((heatmap * 255).astype(np.uint8), cv2.COLORMAP_JET)

  99.         img = cv2.addWeighted(img, 0.6, rgb_heatmap, 0.4, 0)

  100.         return img

  101. 

  102.     def overlay_ignore_mask(self, img, ignore_mask):

  103.         img = img * np.repeat((ignore_mask == 0).astype(np.uint8)[:, :, None], 3, axis=2)

  104.         return img

  105. 

  106.     # -------------------- augment code --------------------------------

  107.     def get_pose_bboxes(self, poses):

  108.         pose_bboxes = []

  109.         for pose in poses:

  110.             x1 = pose[pose[:, 2] > 0][:, 0].min()

  111.             y1 = pose[pose[:, 2] > 0][:, 1].min()

  112.             x2 = pose[pose[:, 2] > 0][:, 0].max()

  113.             y2 = pose[pose[:, 2] > 0][:, 1].max()

  114.             pose_bboxes.append([x1, y1, x2, y2])

  115.         pose_bboxes = np.array(pose_bboxes)

  116.         return pose_bboxes

  117. 

  118.     def resize_data(self, img, ignore_mask, poses, shape):

  119.         """resize img, mask and annotations"""

  120.         img_h, img_w, _ = img.shape

  121. 

  122.         resized_img = cv2.resize(img, shape)

  123.         ignore_mask = cv2.resize(ignore_mask.astype(np.uint8), shape).astype('bool')

  124.         poses[:, :, :2] = (poses[:, :, :2] * np.array(shape) / np.array((img_w, img_h)))

  125.         return resized_img, ignore_mask, poses

  126. 

  127.     def random_resize_img(self, img, ignore_mask, poses):

  128.         h, w, _ = img.shape

  129.         joint_bboxes = self.get_pose_bboxes(poses)

  130.         bbox_sizes = ((joint_bboxes[:, 2:] - joint_bboxes[:, :2] + 1) ** 2).sum(axis=1) ** 0.5

  131. 

  132.         min_scale = params['min_box_size'] / bbox_sizes.min()

  133.         max_scale = params['max_box_size'] / bbox_sizes.max()

  134. 

  135.         min_scale = min(max(min_scale, params['min_scale']), 1)

  136.         max_scale = min(max(max_scale, 1), params['max_scale'])

  137. 

  138.         scale = float((max_scale - min_scale) * random.random() + min_scale)

  139.         shape = (round(w * scale), round(h * scale))

  140. 

  141.         resized_img, resized_mask, resized_poses = self.resize_data(img, ignore_mask, poses, shape)

  142.         return resized_img, resized_mask, resized_poses

  143. 

  144.     def random_rotate_img(self, img, mask, poses):

  145.         h, w, _ = img.shape

  146.         degree = np.random.randn() / 3 * params['max_rotate_degree']

  147.         rad = degree * math.pi / 180

  148.         center = (w / 2, h / 2)

  149.         R = cv2.getRotationMatrix2D(center, degree, 1)

  150.         bbox = (w * abs(math.cos(rad)) + h * abs(math.sin(rad)), w * abs(math.sin(rad)) + h * abs(math.cos(rad)))

  151.         R[0, 2] += bbox[0] / 2 - center[0]

  152.         R[1, 2] += bbox[1] / 2 - center[1]

  153.         rotate_img = cv2.warpAffine(img, R, (int(bbox[0]+0.5), int(bbox[1]+0.5)), flags=cv2.INTER_CUBIC,

  154.                                     borderMode=cv2.BORDER_CONSTANT, borderValue=[127.5, 127.5, 127.5])

  155.         rotate_mask = cv2.warpAffine(mask.astype('uint8')*255, R, (int(bbox[0]+0.5), int(bbox[1]+0.5))) > 0

  156. 

  157.         tmp_poses = np.ones_like(poses)

  158.         tmp_poses[:, :, :2] = poses[:, :, :2].copy()

  159.         tmp_rotate_poses = np.dot(tmp_poses, R.T)  # apply rotation matrix to the poses

  160.         rotate_poses = poses.copy()  # to keep visibility flag

  161.         rotate_poses[:, :, :2] = tmp_rotate_poses

  162.         return rotate_img, rotate_mask, rotate_poses

  163. 

  164.     def random_crop_img(self, img, ignore_mask, poses):

  165.         h, w, _ = img.shape

  166.         insize = self.insize

  167.         joint_bboxes = self.get_pose_bboxes(poses)

  168.         bbox = random.choice(joint_bboxes)  # select a bbox randomly

  169.         bbox_center = bbox[:2] + (bbox[2:] - bbox[:2]) / 2

  170. 

  171.         r_xy = np.random.rand(2)

  172.         perturb = ((r_xy - 0.5) * 2 * params['center_perterb_max'])

  173.         center = (bbox_center + perturb + 0.5).astype('i')

  174. 

  175.         crop_img = np.zeros((insize, insize, 3), 'uint8') + 127.5

  176.         crop_mask = np.zeros((insize, insize), 'bool')

  177. 

  178.         offset = (center - (insize - 1) / 2 + 0.5).astype('i')

  179.         offset_ = (center + (insize - 1) / 2 - (w - 1, h - 1) + 0.5).astype('i')

  180. 

  181.         x1, y1 = (center - (insize-1)/2 + 0.5).astype('i')

  182.         x2, y2 = (center + (insize-1)/2 + 0.5).astype('i')

  183. 

  184.         x1 = max(x1, 0)

  185.         y1 = max(y1, 0)

  186.         x2 = min(x2, w-1)

  187.         y2 = min(y2, h-1)

  188. 

  189.         x_from = -offset[0] if offset[0] < 0 else 0

  190.         y_from = -offset[1] if offset[1] < 0 else 0

  191.         x_to = insize - offset_[0] - 1 if offset_[0] >= 0 else insize - 1

  192.         y_to = insize - offset_[1] - 1 if offset_[1] >= 0 else insize - 1

  193. 

  194.         crop_img[y_from:y_to+1, x_from:x_to+1] = img[y1:y2+1, x1:x2+1].copy()

  195.         crop_mask[y_from:y_to+1, x_from:x_to+1] = ignore_mask[y1:y2+1, x1:x2+1].copy()

  196. 

  197.         poses[:, :, :2] -= offset

  198.         return crop_img.astype('uint8'), crop_mask, poses

  199. 

  200.     def distort_color(self, img):

  201.         img_max = np.broadcast_to(np.array(255, dtype=np.uint8), img.shape[:-1])

  202.         img_min = np.zeros(img.shape[:-1], dtype=np.uint8)

  203. 

  204.         hsv_img = cv2.cvtColor(img.copy(), cv2.COLOR_BGR2HSV).astype(np.int32)

  205.         hsv_img[:, :, 0] = np.maximum(np.minimum(hsv_img[:, :, 0] - 10 + np.random.randint(20 + 1), img_max), img_min) # hue

  206.         hsv_img[:, :, 1] = np.maximum(np.minimum(hsv_img[:, :, 1] - 40 + np.random.randint(80 + 1), img_max), img_min) # saturation

  207.         hsv_img[:, :, 2] = np.maximum(np.minimum(hsv_img[:, :, 2] - 30 + np.random.randint(60 + 1), img_max), img_min) # value

  208.         hsv_img = hsv_img.astype(np.uint8)

  209. 

  210.         distorted_img = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)

  211.         return distorted_img

  212. 

  213.     def flip_img(self, img, mask, poses):

  214.         flipped_img = cv2.flip(img, 1)

  215.         flipped_mask = cv2.flip(mask.astype(np.uint8), 1).astype('bool')

  216.         poses[:, :, 0] = img.shape[1] - 1 - poses[:, :, 0]

  217. 

  218.         def swap_joints(poses, joint_type_, joint_type_2):

  219.             tmp = poses[:, joint_type_].copy()

  220.             poses[:, joint_type_] = poses[:, joint_type_2]

  221.             poses[:, joint_type_2] = tmp

  222. 

  223.         swap_joints(poses, JointType.LeftEye, JointType.RightEye)

  224.         swap_joints(poses, JointType.LeftEar, JointType.RightEar)

  225.         swap_joints(poses, JointType.LeftShoulder, JointType.RightShoulder)

  226.         swap_joints(poses, JointType.LeftElbow, JointType.RightElbow)

  227.         swap_joints(poses, JointType.LeftHand, JointType.RightHand)

  228.         swap_joints(poses, JointType.LeftWaist, JointType.RightWaist)

  229.         swap_joints(poses, JointType.LeftKnee, JointType.RightKnee)

  230.         swap_joints(poses, JointType.LeftFoot, JointType.RightFoot)

  231.         return flipped_img, flipped_mask, poses

  232. 

  233.     def augment_data(self, img, ignore_mask, poses):

  234.         aug_img = img.copy()

  235.         aug_img, ignore_mask, poses = self.random_resize_img(aug_img, ignore_mask, poses)

  236.         aug_img, ignore_mask, poses = self.random_rotate_img(aug_img, ignore_mask, poses)

  237.         aug_img, ignore_mask, poses = self.random_crop_img(aug_img, ignore_mask, poses)

  238.         if np.random.randint(2):

  239.             aug_img = self.distort_color(aug_img)

  240.         if np.random.randint(2):

  241.             aug_img, ignore_mask, poses = self.flip_img(aug_img, ignore_mask, poses)

  242. 

  243.         return aug_img, ignore_mask, poses

  244.     # ------------------------------- end -----------------------------------

  245. 

  246. 

  247.     # ------------------------------ Heatmap ------------------------------------

  248.     # return shape: (height, width)

  249.     def generate_gaussian_heatmap(self, shape, joint, sigma):

  250.         x, y = joint

  251.         grid_x = np.tile(np.arange(shape[1]), (shape[0], 1))

  252.         grid_y = np.tile(np.arange(shape[0]), (shape[1], 1)).transpose()

  253.         grid_distance = (grid_x - x) ** 2 + (grid_y - y) ** 2

  254.         gaussian_heatmap = np.exp(-0.5 * grid_distance / sigma**2)

  255.         return gaussian_heatmap

  256. 

  257.     def generate_heatmaps(self, img, poses, heatmap_sigma):

  258.         heatmaps = np.zeros((0,) + img.shape[:-1])

  259.         sum_heatmap = np.zeros(img.shape[:-1])

  260.         for joint_index in range(len(JointType)):

  261.             heatmap = np.zeros(img.shape[:-1])

  262.             for pose in poses:

  263.                 if pose[joint_index, 2] > 0:

  264.                     jointmap = self.generate_gaussian_heatmap(img.shape[:-1], pose[joint_index][:2], heatmap_sigma)

  265.                     heatmap[jointmap > heatmap] = jointmap[jointmap > heatmap]

  266.                     sum_heatmap[jointmap > sum_heatmap] = jointmap[jointmap > sum_heatmap]

  267.             heatmaps = np.vstack((heatmaps, heatmap.reshape((1,) + heatmap.shape)))

  268.         bg_heatmap = 1 - sum_heatmap  # background channel

  269.         heatmaps = np.vstack((heatmaps, bg_heatmap[None]))

  270.         return heatmaps.astype('f')

  271. 

  272.     def generate_gaussian_heatmap_fast(self, shape, joint, sigma):

  273.         x, y = joint

  274.         grid_x = np.tile(np.arange(shape[1]), (shape[0], 1))

  275.         grid_y = np.tile(np.arange(shape[0]), (shape[1], 1)).transpose()

  276.         grid_x = grid_x + 0.4375

  277.         grid_y = grid_y + 0.4375

  278.         grid_distance = (grid_x - x) ** 2 + (grid_y - y) ** 2

  279.         gaussian_heatmap = np.exp(-0.5 * grid_distance / sigma**2)

  280.         return gaussian_heatmap

  281. 

  282.     def generate_heatmaps_fast(self, img, poses, heatmap_sigma):

  283.         resize_shape = (img.shape[0] // 8, img.shape[1] // 8)

  284.         heatmaps = np.zeros((0,) + resize_shape)

  285.         sum_heatmap = np.zeros(resize_shape)

  286.         for joint_index in range(len(JointType)):

  287.             heatmap = np.zeros(resize_shape)

  288.             for pose in poses:

  289.                 if pose[joint_index, 2] > 0:

  290.                     jointmap = self.generate_gaussian_heatmap_fast(resize_shape, pose[joint_index][:2]/8,

  291.                                                                    heatmap_sigma/8)

  292.                     index_1 = jointmap > heatmap

  293.                     heatmap[index_1] = jointmap[index_1]

  294.                     index_2 = jointmap > sum_heatmap

  295.                     sum_heatmap[index_2] = jointmap[index_2]

  296.             heatmaps = np.vstack((heatmaps, heatmap.reshape((1,) + heatmap.shape)))

  297. 

  298.         bg_heatmap = 1 - sum_heatmap  # background channel

  299.         heatmaps = np.vstack((heatmaps, bg_heatmap[None]))

  300.         return heatmaps.astype('f')

  301.     # ------------------------------ end ------------------------------------

  302. 

  303.     # ------------------------------ PAF ------------------------------------

  304.     # return shape: (2, height, width)

  305.     def generate_constant_paf(self, shape, joint_from, joint_to, paf_width):

  306.         if np.array_equal(joint_from, joint_to): # same joint

  307.             return np.zeros((2,) + shape[:-1])

  308. 

  309.         joint_distance = np.linalg.norm(joint_to - joint_from)

  310.         unit_vector = (joint_to - joint_from) / joint_distance

  311.         rad = np.pi / 2

  312.         # [[0, 1], [-1, 0]]

  313.         rot_matrix = np.array([[np.cos(rad), np.sin(rad)], [-np.sin(rad), np.cos(rad)]])

  314.         # [[u_y], [-u_x]]

  315.         vertical_unit_vector = np.dot(rot_matrix, unit_vector)

  316.         grid_x = np.tile(np.arange(shape[1]), (shape[0], 1))

  317.         grid_y = np.tile(np.arange(shape[0]), (shape[1], 1)).transpose()

  318.         horizontal_inner_product = unit_vector[0] * (grid_x - joint_from[0]) + unit_vector[1] * (grid_y - joint_from[1])

  319.         horizontal_paf_flag = (horizontal_inner_product >= 0) & (horizontal_inner_product <= joint_distance)

  320.         vertical_inner_product = vertical_unit_vector[0] * (grid_x - joint_from[0]) + vertical_unit_vector[1] *\

  321.                                  (grid_y - joint_from[1])

  322.         vertical_paf_flag = np.abs(vertical_inner_product) <= paf_width  # paf_width : 8

  323.         paf_flag = horizontal_paf_flag & vertical_paf_flag

  324.         constant_paf = np.stack((paf_flag, paf_flag)) *\

  325.                        np.broadcast_to(unit_vector, shape[:-1] + (2,)).transpose(2, 0, 1)

  326. 

  327.         return constant_paf

  328. 

  329.     def generate_pafs(self, img, poses, paf_sigma):

  330.         pafs = np.zeros((0,) + img.shape[:-1])

  331. 

  332.         for limb in params['limbs_point']:

  333.             paf = np.zeros((2,) + img.shape[:-1])

  334.             paf_flags = np.zeros(paf.shape) # for constant paf

  335. 

  336.             for pose in poses:

  337.                 joint_from, joint_to = pose[limb]

  338.                 if joint_from[2] > 0 and joint_to[2] > 0:

  339.                     limb_paf = self.generate_constant_paf(img.shape, joint_from[:2], joint_to[:2], paf_sigma) # [2, 368, 368]

  340.                     limb_paf_flags = limb_paf != 0

  341.                     paf_flags += np.broadcast_to(limb_paf_flags[0] | limb_paf_flags[1], limb_paf.shape)

  342. 

  343.                     paf += limb_paf

  344. 

  345.             paf[paf_flags > 0] /= paf_flags[paf_flags > 0]

  346.             pafs = np.vstack((pafs, paf))

  347.         return pafs.astype('f')

  348. 

  349.     def generate_constant_paf_fast(self, shape, joint_from, joint_to, paf_width):

  350.         if np.array_equal(joint_from, joint_to): # same joint

  351.             return np.zeros((2,) + shape[:-1])

  352. 

  353.         joint_distance = np.linalg.norm(joint_to - joint_from)

  354.         unit_vector = (joint_to - joint_from) / joint_distance

  355.         rad = np.pi / 2

  356.         # [[0, 1], [-1, 0]]

  357.         rot_matrix = np.array([[np.cos(rad), np.sin(rad)], [-np.sin(rad), np.cos(rad)]])

  358.         # [[u_y], [-u_x]]

  359.         vertical_unit_vector = np.dot(rot_matrix, unit_vector)

  360.         grid_x = np.tile(np.arange(shape[1]), (shape[0], 1))

  361.         grid_y = np.tile(np.arange(shape[0]), (shape[1], 1)).transpose()

  362.         grid_x = grid_x + 0.4375

  363.         grid_y = grid_y + 0.4375

  364.         horizontal_inner_product = unit_vector[0] * (grid_x - joint_from[0]) + unit_vector[1] * (grid_y - joint_from[1])

  365.         horizontal_paf_flag = (horizontal_inner_product >= 0) & (horizontal_inner_product <= joint_distance)

  366.         vertical_inner_product = vertical_unit_vector[0] * (grid_x - joint_from[0]) + vertical_unit_vector[1] *\

  367.                                  (grid_y - joint_from[1])

  368.         vertical_paf_flag = np.abs(vertical_inner_product) <= paf_width  # paf_width : 8/8 = 1

  369.         paf_flag = horizontal_paf_flag & vertical_paf_flag

  370.         constant_paf = np.stack((paf_flag, paf_flag)) *\

  371.                        np.broadcast_to(unit_vector, shape[:-1] + (2,)).transpose(2, 0, 1)

  372. 

  373.         return constant_paf

  374. 

  375.     def generate_pafs_fast(self, img, poses, paf_sigma):

  376.         resize_shape = (img.shape[0]//8, img.shape[1]//8, 3)

  377.         pafs = np.zeros((0,) + resize_shape[:-1])

  378. 

  379.         for limb in params['limbs_point']:

  380.             paf = np.zeros((2,) + resize_shape[:-1])

  381.             paf_flags = np.zeros(paf.shape) # for constant paf

  382. 

  383.             for pose in poses:

  384.                 joint_from, joint_to = pose[limb]

  385.                 if joint_from[2] > 0 and joint_to[2] > 0:

  386.                     limb_paf = self.generate_constant_paf_fast(resize_shape, joint_from[:2]/8, joint_to[:2]/8, paf_sigma/8) # [2, 368, 368]

  387.                     limb_paf_flags = limb_paf != 0

  388.                     paf_flags += np.broadcast_to(limb_paf_flags[0] | limb_paf_flags[1], limb_paf.shape)

  389. 

  390.                     paf += limb_paf

  391. 

  392.             index_1 = paf_flags > 0

  393.             paf[index_1] /= paf_flags[index_1]

  394.             pafs = np.vstack((pafs, paf))

  395.         return pafs.astype('f')

  396.     # ------------------------------ end ------------------------------------

  397. 

  398.     def get_img_annotation(self, ind=None, img_id=None):

  399.         annotations = None

  400. 

  401.         if ind is not None:

  402.             img_id = self.imgIds[ind]

  403.         anno_ids = self.train.getAnnIds(imgIds=[img_id], iscrowd=None)

  404. 

  405.         # annotation for that image

  406.         if anno_ids:

  407.             annotations_for_img = self.train.loadAnns(anno_ids)

  408. 

  409.             person_cnt = 0

  410.             valid_annotations_for_img = []

  411.             for annotation in annotations_for_img:

  412.                 # if too few keypoints or too small

  413.                 if annotation['num_keypoints'] >= params['min_keypoints'] and annotation['area'] > params['min_area']:

  414.                     person_cnt += 1

  415.                     valid_annotations_for_img.append(annotation)

  416. 

  417.             # if person annotation

  418.             if person_cnt > 0:

  419.                 annotations = valid_annotations_for_img

  420. 

  421.         img_path = os.path.join(self.imgpath, self.train.loadImgs([img_id])[0]['file_name'])

  422.         mask_path = os.path.join(self.maskpath, '{:012d}.png'.format(img_id))

  423.         img = cv2.imread(img_path)

  424.         ignore_mask = cv2.imread(mask_path, 0)

  425.         if ignore_mask is None:

  426.             ignore_mask = np.zeros(img.shape[:2], np.float32)

  427.         else:

  428.             ignore_mask[ignore_mask == 255] = 1

  429. 

  430.         if self.mode == 'eval':

  431.             return img, img_id, annotations_for_img, ignore_mask

  432. 

  433.         return img, img_id, annotations, ignore_mask.astype('f')

  434. 

  435.     def parse_annotation(self, annotations):

  436.         poses = np.zeros((0, len(JointType), 3), dtype=np.int32)

  437. 

  438.         for ann in annotations:

  439.             ann_pose = np.array(ann['keypoints']).reshape(-1, 3)

  440.             pose = np.zeros((1, len(JointType), 3), dtype=np.int32)

  441. 

  442.             # convert poses position

  443.             for i, joint_index in enumerate(params['joint_indices']):

  444.                 pose[0][joint_index] = ann_pose[i]

  445. 

  446.             # compute neck position

  447.             if pose[0][JointType.LeftShoulder][2] > 0 and pose[0][JointType.RightShoulder][2] > 0:

  448.                 pose[0][JointType.Neck][0] = int((pose[0][JointType.LeftShoulder][0] +

  449.                                                   pose[0][JointType.RightShoulder][0]) / 2)

  450.                 pose[0][JointType.Neck][1] = int((pose[0][JointType.LeftShoulder][1] +

  451.                                                   pose[0][JointType.RightShoulder][1]) / 2)

  452.                 pose[0][JointType.Neck][2] = 2

  453. 

  454.             poses = np.vstack((poses, pose))

  455.         return poses

  456. 

  457.     def resize_output(self, input_np, map_h=46, map_w=46):

  458.         if len(input_np.shape) == 3:

  459.             output = np.zeros((input_np.shape[0], map_h, map_w))

  460.             for i in range(input_np.shape[0]):

  461.                 output[i] = cv2.resize(input_np[i], (map_w, map_h))

  462.             return output.astype('f')

  463. 

  464.         input_np = input_np.astype('f')

  465.         output = cv2.resize(input_np, (map_h, map_w))

  466.         return output

  467. 

  468.     def generate_labels(self, img, poses, ignore_mask):

  469.         img, ignore_mask, poses = self.augment_data(img, ignore_mask, poses)

  470.         resized_img, ignore_mask, resized_poses = self.resize_data(img, ignore_mask, poses,

  471.                                                                    shape=(self.insize, self.insize))

  472. 

  473.         resized_heatmaps = self.generate_heatmaps_fast(resized_img, resized_poses, params['heatmap_sigma'])

  474. 

  475.         resized_pafs = self.generate_pafs_fast(resized_img, resized_poses, params['paf_sigma'])

  476. 

  477.         ignore_mask = cv2.morphologyEx(ignore_mask.astype('uint8'), cv2.MORPH_DILATE, np.ones((16, 16))).astype('bool')

  478.         resized_ignore_mask = self.resize_output(ignore_mask)

  479. 

  480. 

  481.         return resized_img, resized_pafs, resized_heatmaps, resized_ignore_mask

  482. 

  483.     def preprocess(self, img):

  484.         x_data = img.astype('f')

  485.         x_data /= 255

  486.         x_data -= 0.5

  487.         x_data = x_data.transpose(2, 0, 1)

  488.         return x_data

  489. 

  490.     def __getitem__(self, i):

  491.         img, img_id, annotations, ignore_mask = self.get_img_annotation(ind=i)

  492. 

  493.         if self.mode in ['eval', 'val']:

  494.             # don't need to make heatmaps/pafs

  495.             return img, np.array([img_id])

  496. 

  497.         # if no annotations are available

  498.         while annotations is None:

  499.             print("none annotations", img_id)

  500.             img_id = self.imgIds[np.random.randint(len(self))]

  501.             img, img_id, annotations, ignore_mask = self.get_img_annotation(img_id=img_id)

  502. 

  503.         poses = self.parse_annotation(annotations)

  504. 

  505.         # TEST

  506.         # return img, poses, ignore_mask

  507. 

  508.         resized_img, pafs, heatmaps, ignore_mask = self.generate_labels(img, poses, ignore_mask)

  509.         resized_img = self.preprocess(resized_img)

  510.         ignore_mask = 1. - ignore_mask

  511. 

  512.         # # TEST

  513.         # print("Shape: ", resized_img.dtype, " ", pafs.dtype, " ", heatmaps.dtype, " ", ignore_mask.dtype)

  514. 

  515.         return resized_img, pafs, heatmaps, ignore_mask

  516. 

  517. 

  518. class DistributedSampler():

  519.     def __init__(self, dataset, rank, group_size, shuffle=True, seed=0):

  520.         self.dataset = dataset

  521.         self.rank = rank

  522.         self.group_size = group_size

  523.         self.dataset_len = len(self.dataset)

  524.         self.num_samplers = int(math.ceil(self.dataset_len * 1.0 / self.group_size))

  525.         self.total_size = self.num_samplers * self.group_size

  526.         self.shuffle = shuffle

  527.         self.seed = seed

  528. 

  529.     def __iter__(self):

  530.         if self.shuffle:

  531.             self.seed = (self.seed + 1) & 0xffffffff

  532.             np.random.seed(self.seed)

  533.             indices = np.random.permutation(self.dataset_len).tolist()

  534.         else:

  535.             indices = list(range(len(self.dataset_len)))

  536.         indices += indices[:(self.total_size - len(indices))]

  537.         indices = indices[self.rank::self.group_size]

  538.         return iter(indices)

  539. 

  540.     def __len__(self):

  541.         return self.num_samplers

  542. 

  543. def valdata(jsonpath, imgpath, rank, group_size, mode='val', maskpath=''):

  544.     #cv2.setNumThreads(0)

  545.     val = ReadJson(jsonpath)

  546.     dataset = txtdataset(val, imgpath, maskpath, params['insize'], mode=mode)

  547.     sampler = DistributedSampler(dataset, rank, group_size)

  548.     ds = de.GeneratorDataset(dataset, ['img', 'img_id'], num_parallel_workers=8, sampler=sampler)

  549.     ds = ds.repeat(1)

  550.     return ds

  551. 

  552. 

  553. def create_dataset(jsonpath, imgpath, maskpath, batch_size, rank, group_size, mode='train', repeat_num=1, shuffle=True,

  554.                    multiprocessing=True, num_worker=20):

  555. 

  556.     train = ReadJson(jsonpath)

  557.     dataset = txtdataset(train, imgpath, maskpath, params['insize'], mode=mode)

  558.     if group_size == 1:

  559.         de_dataset = de.GeneratorDataset(dataset, ["image", "pafs", "heatmaps", "ignore_mask"],

  560.                                          shuffle=shuffle,

  561.                                          num_parallel_workers=num_worker,

  562.                                          python_multiprocessing=multiprocessing)

  563.     else:

  564.         de_dataset = de.GeneratorDataset(dataset, ["image", "pafs", "heatmaps", "ignore_mask"],

  565.                                          shuffle=shuffle,

  566.                                          num_parallel_workers=num_worker,

  567.                                          python_multiprocessing=multiprocessing,

  568.                                          num_shards=group_size,

  569.                                          shard_id=rank)

  570. 

  571.     de_dataset = de_dataset.batch(batch_size=batch_size, drop_remainder=True)

  572.     de_dataset = de_dataset.repeat(repeat_num)

  573. 

  574.     return de_dataset