154 lines
4.7 KiB
Python
154 lines
4.7 KiB
Python
import h5py
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
|
|
|
|
# ----------------------------
|
|
# Helpers: robust mat(v7.3) read
|
|
# ----------------------------
|
|
def _read_array(f, key):
|
|
if key not in f:
|
|
raise KeyError(f"Key '{key}' not found. Available keys: {list(f.keys())}")
|
|
return np.array(f[key])
|
|
|
|
|
|
def _ensure_coords_m2(coords):
|
|
# coords expected: (m,2). Some mat files read as (2,m)
|
|
if coords.ndim != 2:
|
|
raise ValueError(f"coords must be 2D, got {coords.shape}")
|
|
if coords.shape[1] == 2:
|
|
return coords
|
|
if coords.shape[0] == 2:
|
|
return coords.T
|
|
raise ValueError(f"coords should be (m,2) or (2,m), got {coords.shape}")
|
|
|
|
|
|
def _ensure_X_nchw(X, H=512, W=512):
|
|
"""
|
|
Expect geometry:
|
|
(N,1,H,W) OR (H,W,1,N) OR (1,H,W,N)
|
|
"""
|
|
if X.ndim != 4:
|
|
raise ValueError(f"X must be 4D, got {X.shape}")
|
|
|
|
# (N,1,H,W)
|
|
if X.shape[0] > 1 and X.shape[1] == 1 and X.shape[2] == H and X.shape[3] == W:
|
|
return X
|
|
|
|
# (H,W,1,N) -> (N,1,H,W)
|
|
if X.shape[0] == H and X.shape[1] == W and X.shape[2] == 1:
|
|
return np.transpose(X, (3, 2, 0, 1))
|
|
|
|
# (1,H,W,N) -> (N,1,H,W)
|
|
if X.shape[0] == 1 and X.shape[1] == H and X.shape[2] == W:
|
|
return np.transpose(X, (3, 0, 1, 2))
|
|
|
|
raise ValueError(f"Unrecognized X layout: {X.shape}")
|
|
|
|
|
|
def _ensure_Y_Nm(Y, N_expected=None, m_expected=None):
|
|
# expect (N,m) or (m,N)
|
|
if Y.ndim != 2:
|
|
raise ValueError(f"Y must be 2D, got {Y.shape}")
|
|
|
|
if N_expected is not None and Y.shape[0] == N_expected:
|
|
return Y
|
|
if N_expected is not None and Y.shape[1] == N_expected:
|
|
return Y.T
|
|
|
|
if m_expected is not None and Y.shape[1] == m_expected:
|
|
return Y
|
|
if m_expected is not None and Y.shape[0] == m_expected:
|
|
return Y.T
|
|
|
|
return Y
|
|
|
|
|
|
# ----------------------------
|
|
# Dataset (NO normalization here!)
|
|
# ----------------------------
|
|
class MetalensDataset(Dataset):
|
|
def __init__(self, X_geom, coords, Y_re, Y_im):
|
|
"""
|
|
X_geom: (N,1,512,512) already in {0,1} or float
|
|
coords: (m,2) already normalized to [-1,1]
|
|
Y_re: (N,m) already normalized
|
|
Y_im: (N,m) already normalized
|
|
"""
|
|
# cast geometry to float32 for convs
|
|
self.X = torch.tensor(X_geom, dtype=torch.float32)
|
|
self.coords = torch.tensor(coords, dtype=torch.float32) # shared for all samples
|
|
self.Yre = torch.tensor(Y_re, dtype=torch.float32)
|
|
self.Yim = torch.tensor(Y_im, dtype=torch.float32)
|
|
|
|
def __len__(self):
|
|
return self.X.shape[0]
|
|
|
|
def __getitem__(self, idx):
|
|
x = self.X[idx] # (1,512,512)
|
|
y = torch.stack([self.Yre[idx], self.Yim[idx]], dim=-1) # (m,2)
|
|
return x, self.coords, y
|
|
|
|
|
|
# ----------------------------
|
|
# Build datasets
|
|
# ----------------------------
|
|
def build_datasets(
|
|
mat_path="metalens_dataset.mat",
|
|
train_ratio=0.8,
|
|
keys=None,
|
|
):
|
|
"""
|
|
Read your new dataset:
|
|
Rimg : (N,1,512,512) (or equivalent layout)
|
|
coords : (m,2)
|
|
Ere_flat : (N,m)
|
|
Eim_flat : (N,m)
|
|
No norm_stats needed.
|
|
"""
|
|
if keys is None:
|
|
keys = dict(
|
|
X="Rimg",
|
|
coords="coords",
|
|
Yre="Ere_flat",
|
|
Yim="Eim_flat",
|
|
)
|
|
|
|
with h5py.File(mat_path, "r") as f:
|
|
X = _read_array(f, keys["X"])
|
|
coords = _read_array(f, keys["coords"])
|
|
Yre = _read_array(f, keys["Yre"])
|
|
Yim = _read_array(f, keys["Yim"])
|
|
|
|
coords = _ensure_coords_m2(coords)
|
|
X = _ensure_X_nchw(X, 512, 512)
|
|
|
|
N = X.shape[0]
|
|
m = coords.shape[0]
|
|
|
|
Yre = _ensure_Y_Nm(Yre, N_expected=N, m_expected=m)
|
|
Yim = _ensure_Y_Nm(Yim, N_expected=N, m_expected=m)
|
|
|
|
if Yre.shape != (N, m) or Yim.shape != (N, m):
|
|
raise ValueError(f"Y shapes mismatch: Yre={Yre.shape}, Yim={Yim.shape}, expected {(N, m)}")
|
|
|
|
if N < 2:
|
|
raise ValueError(f"N too small: {N}")
|
|
|
|
n_train = int(train_ratio * N)
|
|
n_train = min(max(1, n_train), N - 1)
|
|
|
|
X_train, X_test = X[:n_train], X[n_train:]
|
|
Yre_train, Yre_test = Yre[:n_train], Yre[n_train:]
|
|
Yim_train, Yim_test = Yim[:n_train], Yim[n_train:]
|
|
|
|
train_ds = MetalensDataset(X_train, coords, Yre_train, Yim_train)
|
|
test_ds = MetalensDataset(X_test, coords, Yre_test, Yim_test)
|
|
|
|
print(f"[build_datasets] X: {X.shape}, coords: {coords.shape}, Y: {Yre.shape}")
|
|
print(f"[build_datasets] Split train={len(train_ds)}, test={len(test_ds)}")
|
|
|
|
# norm_stats 不再需要：返回空 dict
|
|
norm_stats = {}
|
|
return train_ds, test_ds, norm_stats |