| 47 | |
| 48 | |
| 49 | class AstroClipCollator: |
| 50 | def __init__( |
| 51 | self, |
| 52 | center_crop: int = 144, |
| 53 | bands: List[str] = ["g", "r", "z"], |
| 54 | m: float = 0.03, |
| 55 | Q: int = 20, |
| 56 | ): |
| 57 | self.center_crop = CenterCrop(center_crop) |
| 58 | self.to_rgb = ToRGB(bands=bands, m=m, Q=Q) |
| 59 | |
| 60 | def _process_images(self, images): |
| 61 | # convert to rgb |
| 62 | img_outs = [] |
| 63 | for img in images: |
| 64 | rgb_img = torch.tensor(self.to_rgb(img)[None, :, :, :]) |
| 65 | img_outs.append(rgb_img) |
| 66 | images = torch.concatenate(img_outs) |
| 67 | |
| 68 | images = self.center_crop(images.permute(0, 3, 2, 1)) |
| 69 | return images |
| 70 | |
| 71 | def __call__(self, samples): |
| 72 | # collate and handle dimensions |
| 73 | samples = default_collate(samples) |
| 74 | # process images |
| 75 | samples["image"] = self._process_images(samples["image"]) |
| 76 | return samples |
no outgoing calls
no test coverage detected