| 684 | return self.im_res / self.pixel_count |
| 685 | |
| 686 | class ImageSpliterTh: |
| 687 | def __init__(self, im, pch_size, stride, sf=1): |
| 688 | ''' |
| 689 | Input: |
| 690 | im: n x c x h x w, torch tensor, float, low-resolution image in SR |
| 691 | pch_size, stride: patch setting |
| 692 | sf: scale factor in image super-resolution |
| 693 | ''' |
| 694 | assert stride <= pch_size |
| 695 | self.stride = stride |
| 696 | self.pch_size = pch_size |
| 697 | self.sf = sf |
| 698 | |
| 699 | bs, chn, height, width= im.shape |
| 700 | self.height_starts_list = self.extract_starts(height) |
| 701 | self.width_starts_list = self.extract_starts(width) |
| 702 | self.length = self.__len__() |
| 703 | self.num_pchs = 0 |
| 704 | |
| 705 | self.im_ori = im |
| 706 | self.im_res = torch.zeros([bs, chn, height*sf, width*sf], dtype=im.dtype, device=im.device) |
| 707 | self.pixel_count = torch.zeros([bs, chn, height*sf, width*sf], dtype=im.dtype, device=im.device) |
| 708 | self.weight = self._gaussian_weights(pch_size, pch_size, bs, im.device) |
| 709 | |
| 710 | def extract_starts(self, length): |
| 711 | if length <= self.pch_size: |
| 712 | starts = [0,] |
| 713 | else: |
| 714 | starts = list(range(0, length, self.stride)) |
| 715 | for i in range(len(starts)): |
| 716 | if starts[i] + self.pch_size > length: |
| 717 | starts[i] = length - self.pch_size |
| 718 | starts = sorted(set(starts), key=starts.index) |
| 719 | return starts |
| 720 | |
| 721 | def __len__(self): |
| 722 | return len(self.height_starts_list) * len(self.width_starts_list) |
| 723 | |
| 724 | def __iter__(self): |
| 725 | return self |
| 726 | |
| 727 | def __next__(self): |
| 728 | if self.num_pchs < self.length: |
| 729 | w_start_idx = self.num_pchs // len(self.height_starts_list) |
| 730 | w_start = self.width_starts_list[w_start_idx] |
| 731 | w_end = w_start + self.pch_size |
| 732 | |
| 733 | h_start_idx = self.num_pchs % len(self.height_starts_list) |
| 734 | h_start = self.height_starts_list[h_start_idx] |
| 735 | h_end = h_start + self.pch_size |
| 736 | |
| 737 | pch = self.im_ori[:, :, h_start:h_end, w_start:w_end,] |
| 738 | |
| 739 | h_start *= self.sf |
| 740 | h_end *= self.sf |
| 741 | w_start *= self.sf |
| 742 | w_end *= self.sf |
| 743 | |