| """Synth4: Synth3 synthetic pairs + real EarthView single-date 'no change' pairs (fraction P_EV).""" |
| import os |
| import time |
|
|
| import cv2 |
| import numpy as np |
| import torch |
|
|
| from train2 import CROP |
| from train3 import Synth3, degrade |
|
|
|
|
| class Synth4(Synth3): |
| def __init__(self, n, p_ev=0.2): |
| super().__init__(n) |
| self.p_ev = p_ev |
| self.ev = np.load("/workspace/prep/ev_imgs.npy", mmap_mode="r") |
| print("earthview singles", self.ev.shape, flush=True) |
|
|
| def __getitem__(self, idx): |
| rng = np.random.default_rng((idx * 7919 + os.getpid() * 104729 + time.time_ns()) % 2**32) |
| if rng.random() >= self.p_ev: |
| return super().__getitem__(idx) |
| img = self.ev[rng.integers(len(self.ev))] |
| size = int(CROP * rng.uniform(0.55, 0.7)) |
| r, c = rng.integers(0, img.shape[0] - size + 1, 2) |
| crop = cv2.resize(np.ascontiguousarray(img[r:r + size, c:c + size]), (CROP, CROP), interpolation=cv2.INTER_CUBIC) |
| pre, post = degrade(self.photo(crop, rng), rng), degrade(self.photo(crop.copy(), rng), rng) |
| if rng.random() < 0.8: |
| M = cv2.getRotationMatrix2D((CROP / 2, CROP / 2), rng.uniform(-1.5, 1.5), rng.uniform(0.98, 1.02)) |
| M[:, 2] += rng.uniform(-4, 4, 2) |
| pre = cv2.warpAffine(pre, M, (CROP, CROP), borderMode=cv2.BORDER_REFLECT) |
| k = rng.integers(8) |
| def d4(a): |
| a = np.rot90(a, k % 4) |
| return np.ascontiguousarray(a[:, ::-1] if k >= 4 else a) |
| pre, post = d4(pre), d4(post) |
| ig = np.full((CROP, CROP), 255, np.int64) |
| t = lambda a: torch.from_numpy(a.transpose(2, 0, 1).copy()) |
| return (t(pre), t(post), torch.zeros(2, CROP, CROP), torch.zeros(2), torch.from_numpy(ig), torch.from_numpy(ig.copy())) |
|
|
