| """Synth5: real JL1-CD pairs (pseudo class labels) with prob P_REAL, else Synth3 synthetic pairs.""" |
| import os, time |
| import cv2, numpy as np, torch |
| from train2 import CROP |
| from train3 import Synth3 |
|
|
| class Synth5(Synth3): |
| def __init__(self, n, p_real=0.5): |
| super().__init__(n) |
| self.p_real = p_real |
| self.jpre = np.load("/workspace/prep/jl_pre.npy", mmap_mode="r"); self.jpost = np.load("/workspace/prep/jl_post.npy", mmap_mode="r") |
| self.jb = np.load("/workspace/prep/jl_lb.npy", mmap_mode="r"); self.jt = np.load("/workspace/prep/jl_lt.npy", mmap_mode="r") |
| print("JL1 real pairs", self.jpre.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_real: |
| return super().__getitem__(idx) |
| i = rng.integers(len(self.jpre)); size = int(CROP * rng.uniform(0.85, 1.15)); r, c = rng.integers(0, 512 - size + 1, 2) |
| cr = lambda a, it: cv2.resize(np.ascontiguousarray(a[i, r:r + size, c:c + size]), (CROP, CROP), interpolation=it) |
| pre, post = cr(self.jpre, cv2.INTER_AREA), cr(self.jpost, cv2.INTER_AREA) |
| lb, lt = cr(self.jb, cv2.INTER_NEAREST), cr(self.jt, cv2.INTER_NEAREST) |
| if rng.random() < 0.3: pre = self.photo(pre, rng) |
| if rng.random() < 0.3: post = self.photo(post, rng) |
| 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, lb, lt = d4(pre), d4(post), d4(lb), d4(lt) |
| ig = np.full((CROP, CROP), 255, np.int64) |
| t = lambda a: torch.from_numpy(a.transpose(2, 0, 1).copy()) |
| lab = np.stack([lb, lt]).astype(np.float32); pres = np.array([lb.sum() >= 20, lt.sum() >= 20], np.float32) |
| return (t(pre), t(post), torch.from_numpy(lab), torch.from_numpy(pres), torch.from_numpy(ig), torch.from_numpy(ig.copy())) |
|
|
