import sys sys.path.append('./src') import matplotlib matplotlib.use('Agg') from pathlib import Path from omegaconf import OmegaConf import torch from tqdm import tqdm from matplotlib import pyplot as plt import numpy as np import h5py from models import get_model_class from datasets import singlePhotonDataset torch.manual_seed(0) torch.cuda.manual_seed(0) np.random.seed(0) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False conf = OmegaConf.load("Configs/infer_1photon.yaml") NX, NY = conf.data.names[conf.experiment.name].NX, conf.data.names[conf.experiment.name].NY def inv0(p): return p def inv1(p): return torch.stack([-p[..., 0], p[..., 1]], dim=-1) def inv2(p): return torch.stack([p[..., 0], -p[..., 1]], dim=-1) def inv3(p): return -p def inv4(p): return torch.stack([p[..., 1], p[..., 0]], dim=-1) def inv5(p): return torch.stack([p[..., 1], -p[..., 0]], dim=-1) def inv6(p): return torch.stack([-p[..., 1], p[..., 0]], dim=-1) def inv7(p): return torch.stack([-p[..., 1], -p[..., 0]], dim=-1) INVERSE_TRANSFORMS = { 0: inv0, 1: inv1, 2: inv2, 3: inv3, 4: inv4, 5: inv5, 6: inv6, 7: inv7, } def apply_inverse_transforms(predictions: torch.Tensor, numberOfAugOps: int) -> torch.Tensor: N = predictions.shape[0] // numberOfAugOps preds = predictions.view(N, numberOfAugOps, 2) corrected = torch.zeros_like(preds) for idx in range(numberOfAugOps): corrected[:, idx, :] = INVERSE_TRANSFORMS[idx](preds[:, idx, :]) return corrected.mean(dim=1) def prepare_output_folder(conf): if conf.data.normalize: normalize_suffix = '_normalized' else: normalize_suffix = '' output_base = Path(conf.experiment.output_base) / conf.experiment.name / conf.model.experiment_name / f'augX{conf.inference.num_aug_ops}{normalize_suffix}' output_base.mkdir(parents=True, exist_ok=True) OmegaConf.save(conf, output_base / 'config.yaml') return output_base def get_files_list(conf): task_conf = conf.data.names[conf.experiment.name] file_pattern = task_conf.file_pattern start, end = task_conf.file_range files = [str(Path(conf.data.sample_folder) / file_pattern.format(i)) for i in range(start, end)] return files def run_inference(model, data_loader, conf): all_predictions = [] all_reference_points = data_loader.dataset.referencePoint with torch.no_grad(): for batch in tqdm(data_loader, desc="Inferring"): inputs, _ = batch inputs = inputs.to('cuda') outputs = model(inputs)[:, :2].cpu() # 只取 x, y all_predictions.append(outputs) all_predictions = torch.cat(all_predictions, dim=0) all_predictions = apply_inverse_transforms(all_predictions, conf.inference.num_aug_ops) offset = [inputs.shape[-2] / 2., inputs.shape[-1] / 2.] offset = torch.tensor(offset).unsqueeze(0) # (1, 2) all_predictions += offset print(f'mean x = {torch.mean(all_predictions[:, 0]):.4f}, std x = {torch.std(all_predictions[:, 0]):.4f}') print(f'mean y = {torch.mean(all_predictions[:, 1]):.4f}, std y = {torch.std(all_predictions[:, 1]):.4f}') return all_predictions.numpy(), all_reference_points def accumulate_hits(predictions: np.ndarray, reference_points: np.ndarray, binning_factor: int): ml_super_frame = np.zeros((NY * binning_factor, NX * binning_factor)) count_frame = np.zeros((NY, NX)) subpixel_dist = np.zeros((binning_factor, binning_factor)) # 绝对坐标 = 预测亚像素 + 参考点(像素左下角) absolute_positions = predictions + reference_points[:, :2] # 超分辨帧 (binning) hit_x = np.floor(absolute_positions[:, 0] * binning_factor).astype(int) hit_y = np.floor(absolute_positions[:, 1] * binning_factor).astype(int) np.add.at(ml_super_frame, (hit_y, hit_x), 1) # 计数帧 (按参考点像素索引) ref_x = (reference_points[:, 0] + 1).astype(int) # 参考点是左下角，+1 得像素索引 ref_y = (reference_points[:, 1] + 1).astype(int) np.add.at(count_frame, (ref_y, ref_x), 1) # 亚像素分布 sub_x = np.floor((absolute_positions[:, 0] % 1) * binning_factor).astype(int) sub_y = np.floor((absolute_positions[:, 1] % 1) * binning_factor).astype(int) np.add.at(subpixel_dist, (sub_y, sub_x), 1) return ml_super_frame, count_frame, subpixel_dist def save_results(ml_super_frame, count_frame, subpixel_dist, roi: list, binning_factor: int, output_dir: Path): x_min, x_max, y_min, y_max = roi # 1. super-resolution frame plt.figure(figsize=(8, 8)) plt.imshow(ml_super_frame[y_min*binning_factor:y_max*binning_factor, x_min*binning_factor:x_max*binning_factor], origin='lower', extent=[x_min, x_max, y_min, y_max]) plt.colorbar(label='Counts') plt.title('ML Super-Resolution Frame') plt.xlabel('X (pixel)') plt.ylabel('Y (pixel)') plt.savefig(output_dir / '1Photon_ML_superFrame.png', dpi=300, bbox_inches='tight') plt.close() np.save(output_dir / '1Photon_ML_superFrame.npy', ml_super_frame) # 2. count frame plt.figure(figsize=(8, 8)) plt.imshow(count_frame[y_min:y_max, x_min:x_max], origin='lower', extent=[x_min, x_max, y_min, y_max]) plt.colorbar(label='Counts') plt.title('Photon Count Frame') plt.xlabel('X (pixel)') plt.ylabel('Y (pixel)') plt.savefig(output_dir / '1Photon_count_Frame.png', dpi=300, bbox_inches='tight') plt.close() np.save(output_dir / '1Photon_count_Frame.npy', count_frame) # 3. sub-pixel distribution plt.figure(figsize=(8, 8)) plt.imshow(subpixel_dist, origin='lower', extent=[0, binning_factor, 0, binning_factor], cmap='viridis') plt.colorbar(label='Counts') plt.title('Sub-pixel Distribution') plt.xlabel(f'Sub-pixel X (1/{binning_factor} pixel)') plt.ylabel(f'Sub-pixel Y (1/{binning_factor} pixel)') plt.savefig(output_dir / '1Photon_subpixel_Distribution.png', dpi=300, bbox_inches='tight') plt.close() np.save(output_dir / '1Photon_subpixel_Distribution.npy', subpixel_dist) rms, mean = np.std(subpixel_dist), np.mean(subpixel_dist) print(f"[Plotting]: Sub-pixel distribution: RMS/Mean: {rms/mean:.4f}, expected value = {1/np.sqrt(mean):.4f} for uniform distribution") print(f"Results saved to: {output_dir}") if __name__ == "__main__": ### output preparation output_dir = prepare_output_folder(conf) ### model loading model_version = conf.model.experiment_name.split('_v')[-1][:6] model = get_model_class(model_version)().cuda() model.load_state_dict(torch.load(f'{conf.model.base_dir}/{conf.model.experiment_name}/Models/{conf.model.name}', weights_only=True)) model.eval() ### data loading files_list = get_files_list(conf) roi = conf.data.names[conf.experiment.name].roi BinningFactor = conf.inference.binning_factor numberOfAugOps = conf.inference.num_aug_ops flag_normalize = conf.data.normalize nChunks = int(np.ceil(len(files_list) / conf.inference.chunk_size)) ml_super_frame = np.zeros((NY * BinningFactor, NX * BinningFactor)) count_frame = np.zeros((NY, NX)) subpixel_dist = np.zeros((BinningFactor, BinningFactor)) ### Inference loop for idxChunk in range(nChunks): start_idx = idxChunk * conf.inference.chunk_size end_idx = min(start_idx + conf.inference.chunk_size, len(files_list)) chunk_files = files_list[start_idx:end_idx] print(f'[Inferring] Chunk {idxChunk+1}/{nChunks}: Loading files {start_idx} to {end_idx}...') dataset = singlePhotonDataset( chunk_files, sampleRatio=1.0, datasetName='Inference', numberOfAugOps=conf.inference.num_aug_ops, normalize=conf.data.normalize ) data_loader = torch.utils.data.DataLoader( dataset, batch_size=conf.data.batch_size, shuffle=False, num_workers=conf.data.num_workers, pin_memory=True ) predictions, ref_points = run_inference(model, data_loader, conf) chunk_super, chunk_count, chunk_subpixel = accumulate_hits( predictions, ref_points, binning_factor=BinningFactor ) ml_super_frame += chunk_super count_frame += chunk_count subpixel_dist += chunk_subpixel save_results(ml_super_frame, count_frame, subpixel_dist, roi=roi, binning_factor=BinningFactor, output_dir=output_dir)