(self, raw_img, fname, specific_info=None, return_mask=False)
| 958 | return [crop_top, crop_bottom] |
| 959 | |
| 960 | def __call__(self, raw_img, fname, specific_info=None, return_mask=False): |
| 961 | h, w, c = raw_img.shape |
| 962 | output = {} |
| 963 | |
| 964 | if specific_info is not None: |
| 965 | # debug-mode |
| 966 | head, neck, waist, foot = specific_info |
| 967 | else: |
| 968 | if fname not in self.info.keys(): |
| 969 | print("split_info not found".center(20, '=')) |
| 970 | print(fname) |
| 971 | return output |
| 972 | head, neck, waist, foot = self.info[fname] |
| 973 | |
| 974 | src_point = np.array([ |
| 975 | [0, head * h], [w, head * h], |
| 976 | [0, neck * h], [w, neck * h], |
| 977 | [0, waist * h], [w, waist * h], |
| 978 | [0, foot * h], [w, foot * h], |
| 979 | ], dtype=np.float32) |
| 980 | |
| 981 | affined_img, trans_matrix = self.affine_image(raw_img, src_point, self.template_point, |
| 982 | shape=self.template_shape) |
| 983 | |
| 984 | if np.random.rand() <= self.crop_prob: |
| 985 | crop_region = self._get_crop_region(affined_img, aug_type=self.aug_type) |
| 986 | else: |
| 987 | crop_region = [-1, -1] |
| 988 | if 0 < crop_region[0] < crop_region[1]: |
| 989 | affined_img[crop_region[0]: crop_region[1], ...] = 0 |
| 990 | |
| 991 | output['affined_image'] = affined_img |
| 992 | # 'theta': theta_tensor |
| 993 | |
| 994 | if return_mask: |
| 995 | # do align&crop as raw_img |
| 996 | fake_img = np.full(raw_img.shape, 125).astype(np.uint8) |
| 997 | affined_fake_img, _ = self.affine_image(fake_img, src_point, self.template_point, |
| 998 | shape=self.template_shape) |
| 999 | if len(affined_fake_img.shape) == 3: |
| 1000 | lum = np.max(affined_fake_img, axis=2) |
| 1001 | else: |
| 1002 | lum = affined_fake_img |
| 1003 | binary_mask = np.ones_like(lum) |
| 1004 | binary_mask = np.where(lum == 0, 0, binary_mask) |
| 1005 | if 0 < crop_region[0] < crop_region[1]: |
| 1006 | binary_mask[crop_region[0]: crop_region[1]:, ...] = 0 |
| 1007 | output['binary_mask'] = binary_mask |
| 1008 | |
| 1009 | return output |
nothing calls this directly
no test coverage detected