| 561 | return view_samples, ray_orgs, ray_dirs |
| 562 | |
| 563 | def get_rays(self, idx): |
| 564 | idxs = self.shuffle_idxs[self.samples_per_view * idx:self.samples_per_view * (idx+1)] |
| 565 | ray_dirs = self.all_ray_dirs.view(-1, 3)[idxs, ...] |
| 566 | ray_orgs = self.all_ray_orgs.view(-1, 3)[idxs, ...] |
| 567 | |
| 568 | if self.multiscale: |
| 569 | view_samples = [mimg.view(-1, 3)[idxs] for mimg in self.multiscale_imgs] |
| 570 | |
| 571 | if self.supervise_hr: |
| 572 | view_samples = [view_samples[-1] for _ in view_samples] |
| 573 | else: |
| 574 | img = self.all_imgs.view(-1, 3)[idxs, ...] |
| 575 | view_samples = img.reshape(-1, 3) |
| 576 | |
| 577 | self.hit[idxs] += 1 |
| 578 | |
| 579 | return view_samples, ray_orgs, ray_dirs |
| 580 | |
| 581 | def __getitem__(self, idx): |
| 582 | |