(self, raw_img, fname, specific_info=None, return_mask=False)
| 654 | return [crop_top, crop_bottom] |
| 655 | |
| 656 | def __call__(self, raw_img, fname, specific_info=None, return_mask=False): |
| 657 | h, w, c = raw_img.shape |
| 658 | output = {} |
| 659 | |
| 660 | if specific_info is not None: |
| 661 | # debug-mode |
| 662 | head, neck, waist, foot = specific_info |
| 663 | else: |
| 664 | if fname not in self.info.keys(): |
| 665 | print("split_info not found".center(20, '=')) |
| 666 | print(fname) |
| 667 | return output |
| 668 | head, neck, waist, foot = self.info[fname] |
| 669 | |
| 670 | src_point = np.array([ |
| 671 | [0, head * h], [w, head * h], |
| 672 | [0, neck * h], [w, neck * h], |
| 673 | [0, waist * h], [w, waist * h], |
| 674 | [0, foot * h], [w, foot * h], |
| 675 | ], dtype=np.float32) |
| 676 | |
| 677 | affined_img, trans_matrix = self.affine_image(raw_img, src_point, self.template_point, |
| 678 | shape=self.template_shape) |
| 679 | |
| 680 | if np.random.rand() <= self.crop_prob: |
| 681 | crop_region = self._get_crop_region(affined_img, aug_type=self.aug_type) |
| 682 | else: |
| 683 | crop_region = [-1, -1] |
| 684 | if 0 < crop_region[0] < crop_region[1]: |
| 685 | affined_img[crop_region[0]: crop_region[1], ...] = 0 |
| 686 | |
| 687 | output['affined_image'] = affined_img |
| 688 | # 'theta': theta_tensor |
| 689 | |
| 690 | if return_mask: |
| 691 | # do align&crop as raw_img |
| 692 | fake_img = np.full(raw_img.shape, 125).astype(np.uint8) |
| 693 | affined_fake_img, _ = self.affine_image(fake_img, src_point, self.template_point, |
| 694 | shape=self.template_shape) |
| 695 | if len(affined_fake_img.shape) == 3: |
| 696 | lum = np.max(affined_fake_img, axis=2) |
| 697 | else: |
| 698 | lum = affined_fake_img |
| 699 | binary_mask = np.ones_like(lum) |
| 700 | binary_mask = np.where(lum == 0, 0, binary_mask) |
| 701 | if 0 < crop_region[0] < crop_region[1]: |
| 702 | binary_mask[crop_region[0]: crop_region[1]:, ...] = 0 |
| 703 | output['binary_mask'] = binary_mask |
| 704 | |
| 705 | return output |
nothing calls this directly
no test coverage detected