diff --git a/Infer_1Photon.py b/Infer_1Photon.py index c4fc9a8..90f4323 100644 --- a/Infer_1Photon.py +++ b/Infer_1Photon.py @@ -81,32 +81,36 @@ def run_inference(model, data_loader, conf): def accumulate_hits(predictions: np.ndarray, reference_points: np.ndarray, binning_factor: int, number_of_subframes: int = 20): - ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor)) - count_frame = np.zeros((NY, NX)) - subpixel_dist = np.zeros((binning_factor, binning_factor)) + ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor), dtype=np.uint32) + count_frame = np.zeros((NY, NX), dtype=np.uint32) + subpixel_dist = np.zeros((binning_factor, binning_factor), dtype=np.uint32) - # absolute coordinate = predicted subpixel + reference point - absolute_positions = predictions + reference_points[:, :2] - - # super resolution frames (binning) - hit_x = np.floor(absolute_positions[:, 0] * binning_factor).astype(int) - hit_y = np.floor(absolute_positions[:, 1] * binning_factor).astype(int) + chunk_size = len(predictions) // number_of_subframes for i in range(number_of_subframes): - start_idx = i * len(predictions) // number_of_subframes - end_idx = min((i + 1) * len(predictions) // number_of_subframes, len(predictions)) - np.add.at(ml_super_frames[i], (hit_y[start_idx:end_idx], hit_x[start_idx:end_idx]), 1) - - # count frame (by reference point pixel index) - ref_x = (reference_points[:, 0] + 1).astype(int) # reference point is lower-left corner, +1 to get pixel index - ref_y = (reference_points[:, 1] + 1).astype(int) - np.add.at(count_frame, (ref_y, ref_x), 1) - - # subpixel distribution - 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) - + start_idx = i * chunk_size + end_idx = (i + 1) * chunk_size if i < number_of_subframes - 1 else len(predictions) + + pred_chunk = predictions[start_idx:end_idx] + ref_chunk = reference_points[start_idx:end_idx] + + abs_pos_chunk = pred_chunk + ref_chunk[:, :2] + + # --- Super resolution frames --- + hit_x = np.floor(abs_pos_chunk[:, 0] * binning_factor).astype(np.int32) + hit_y = np.floor(abs_pos_chunk[:, 1] * binning_factor).astype(np.int32) + np.add.at(ml_super_frames[i], (hit_y, hit_x), 1) + + # --- Count frame --- + ref_x = (ref_chunk[:, 0] + 1).astype(np.int32) + ref_y = (ref_chunk[:, 1] + 1).astype(np.int32) + np.add.at(count_frame, (ref_y, ref_x), 1) + + # --- Subpixel distribution --- + sub_x = np.floor((abs_pos_chunk[:, 0] % 1) * binning_factor).astype(np.int32) + sub_y = np.floor((abs_pos_chunk[:, 1] % 1) * binning_factor).astype(np.int32) + np.add.at(subpixel_dist, (sub_y, sub_x), 1) + return ml_super_frames, count_frame, subpixel_dist def save_results(ml_super_frames, ml_super_frames_eta, diff --git a/Infer_3Photon.py b/Infer_3Photon.py index 919226f..8f8e036 100644 --- a/Infer_3Photon.py +++ b/Infer_3Photon.py @@ -3,6 +3,8 @@ sys.path.append('./src') import matplotlib matplotlib.use('Agg') +from etaInterpolationFromPoints import interpolate_eta_from_points + from pathlib import Path from omegaconf import OmegaConf import torch @@ -47,11 +49,7 @@ def apply_inverse_transforms(predictions: torch.Tensor, numberOfAugOps: int) -> 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 = Path(conf.experiment.output_base) / conf.experiment.name / (f'3ph_{conf.model.experiment_name}') / (f'augX{conf.inference.num_aug_ops}') output_base.mkdir(parents=True, exist_ok=True) OmegaConf.save(conf, output_base / 'config.yaml') return output_base @@ -75,45 +73,52 @@ def run_inference(model, data_loader, conf): all_predictions = torch.cat(all_predictions, dim=0) # all_predictions = apply_inverse_transforms(all_predictions, conf.inference.num_aug_ops) all_predictions += torch.tensor([NSIZE/2., NSIZE/2.]).unsqueeze(0) # adjust back to original coordinate system - print(f'mean x = {torch.mean(all_predictions[:, 0])}, std x = {torch.std(all_predictions[:, 0])}') - print(f'mean y = {torch.mean(all_predictions[:, 1])}, std y = {torch.std(all_predictions[:, 1])}') + print(f'[Inference]: mean x = {torch.mean(all_predictions[:, 0])}, std x = {torch.std(all_predictions[:, 0])}') + print(f'[Inference]: mean y = {torch.mean(all_predictions[:, 1])}, std y = {torch.std(all_predictions[:, 1])}') referencePoints = data_loader.dataset.referencePoint ### the lower-left corner of the cluster in absolute coordinate referencePoints = np.repeat(referencePoints, 3, axis=0) ### duplicate reference points for 3-photon clusters return all_predictions.numpy(), referencePoints def accumulate_hits(predictions: np.ndarray, reference_points: np.ndarray, - binning_factor: int): - ### ret - ml_super_frame = np.zeros((NY*binning_factor, NX*binning_factor), dtype=np.int32) + binning_factor: int, number_of_subframes: int = 20): + ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor)) count_frame = np.zeros((NY, NX), dtype=np.int32) subpixel_dist = np.zeros((binning_factor, binning_factor), dtype=np.int32) ### absolute coordinate = predicted subpixel + reference point absolute_positions = predictions + reference_points - hit_x_superpixel_idx = np.floor(absolute_positions[:, 0] * binning_factor).astype(int) - hit_x_superpixel_idx = np.clip(hit_x_superpixel_idx, 0, NX*binning_factor-1) - hit_y_superpixel_idx = np.floor(absolute_positions[:, 1] * binning_factor).astype(int) - hit_y_superpixel_idx = np.clip(hit_y_superpixel_idx, 0, NY*binning_factor-1) - np.add.at(ml_super_frame, (hit_y_superpixel_idx, hit_x_superpixel_idx), 1) - hit_x_pixel_idx = np.floor(absolute_positions[:, 0]).astype(int) - hit_x_pixel_idx = np.clip(hit_x_pixel_idx, 0, NX-1) - hit_y_pixel_idx = np.floor(absolute_positions[:, 1]).astype(int) - hit_y_pixel_idx = np.clip(hit_y_pixel_idx, 0, NY-1) - np.add.at(count_frame, (hit_y_pixel_idx, hit_x_pixel_idx), 1) + # super resolution frames (binning) + hit_x = np.floor(absolute_positions[:, 0] * binning_factor).astype(int) + hit_y = np.floor(absolute_positions[:, 1] * binning_factor).astype(int) + + for i in range(number_of_subframes): + start_idx = i * len(predictions) // number_of_subframes + end_idx = min((i + 1) * len(predictions) // number_of_subframes, len(predictions)) + np.add.at(ml_super_frames[i], (hit_y[start_idx:end_idx], hit_x[start_idx:end_idx]), 1) + + # count frame (by reference point pixel index) + ref_x = (reference_points[:, 0] + 1).astype(int) # reference point is lower-left corner, +1 to get pixel index + ref_y = (reference_points[:, 1] + 1).astype(int) + np.add.at(count_frame, (ref_y, ref_x), 1) + + # subpixel distribution + 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_frames, count_frame, subpixel_dist - subpixel_x_idx = np.floor((absolute_positions[:, 0] % 1) * binning_factor).astype(int) - subpixel_y_idx = np.floor((absolute_positions[:, 1] % 1) * binning_factor).astype(int) - np.add.at(subpixel_dist, (subpixel_y_idx, subpixel_x_idx), 1) - - return ml_super_frame, count_frame, subpixel_dist - -def save_results(ml_super_frame, count_frame, subpixel_dist, +def save_results(ml_super_frames, ml_super_frames_eta, + count_frame, subpixel_dist, subpixel_dist_eta, roi: list, binning_factor: int, output_dir: Path): x_st, x_ed, y_st, y_ed = roi # 1. super-resolution frame + np.save(output_dir / '3Photon_ML_superFrames.npy', ml_super_frames) + + ml_super_frame = np.sum(ml_super_frames, axis=0) plt.figure(figsize=(8, 8)) plt.imshow(ml_super_frame[y_st*binning_factor:y_ed*binning_factor, x_st*binning_factor:x_ed*binning_factor], origin='lower', extent=[x_st, x_ed, y_st, y_ed]) plt.colorbar(label='Counts') @@ -124,7 +129,21 @@ def save_results(ml_super_frame, count_frame, subpixel_dist, plt.clf() np.save(output_dir / '3Photon_ML_superFrame.npy', ml_super_frame) - # 2. count frame + # 2. super-resolution frame with eta interpolation + np.save(output_dir / '3Photon_ML_superFrames_etaInterpolated.npy', ml_super_frames_eta) + + ml_super_frame_eta = np.sum(ml_super_frames_eta, axis=0) + plt.figure(figsize=(8, 8)) + plt.imshow(ml_super_frame_eta[y_st*binning_factor:y_ed*binning_factor, x_st*binning_factor:x_ed*binning_factor], origin='lower', extent=[x_st, x_ed, y_st, y_ed]) + plt.colorbar(label='Counts') + plt.title('ML Super-Resolution Frame with Eta Interpolation') + plt.xlabel('X (pixel)') + plt.ylabel('Y (pixel)') + plt.savefig(output_dir / '3Photon_ML_superFrame_etaInterpolated.png', dpi=300, bbox_inches='tight') + plt.clf() + np.save(output_dir / '3Photon_ML_superFrame_etaInterpolated.npy', ml_super_frame_eta) + + # 3. count frame plt.imshow(count_frame[y_st:y_ed, x_st:x_ed], origin='lower', extent=[x_st, x_ed, y_st, y_ed]) plt.colorbar(label='Counts') plt.title('Photon Count Frame') @@ -134,7 +153,7 @@ def save_results(ml_super_frame, count_frame, subpixel_dist, plt.clf() np.save(output_dir / '3Photon_count_Frame.npy', count_frame) - # 3. subpixel distribution + # 4. subpixel distribution plt.imshow(subpixel_dist, origin='lower', extent=[0, 1, 0, 1]) plt.colorbar(label='Counts') plt.title('Subpixel Distribution') @@ -146,7 +165,19 @@ def save_results(ml_super_frame, count_frame, subpixel_dist, std, mean = np.std(subpixel_dist), np.mean(subpixel_dist) print(f"[Plotting]: Sub-pixel distribution: RMS/Mean: {std/mean:.4f}, expected value = {1/np.sqrt(mean):.4f} for uniform distribution") - print(f"Results saved to: {output_dir}") + # 5. subpixel distribution with eta interpolation + plt.imshow(subpixel_dist_eta, origin='lower', extent=[0, 1, 0, 1]) + plt.colorbar(label='Counts') + plt.title('Subpixel Distribution with Eta Interpolation') + plt.xlabel('Subpixel X') + plt.ylabel('Subpixel Y') + plt.savefig(output_dir / '3Photon_subpixel_Distribution_etaInterpolated.png', dpi=300, bbox_inches='tight') + plt.close() + np.save(output_dir / '3Photon_subpixel_Distribution_etaInterpolated.npy', subpixel_dist_eta) + std_eta, mean_eta = np.std(subpixel_dist_eta), np.mean(subpixel_dist_eta) + print(f"[Plotting]: Sub-pixel distribution with eta interpolation: RMS/Mean: {std_eta/mean_eta:.4f}, expected value = {1/np.sqrt(mean_eta):.4f} for uniform distribution") + + print(f"[Plotting]: Results saved to: {output_dir}") if __name__ == "__main__": ### output folder preparation @@ -166,6 +197,9 @@ if __name__ == "__main__": flag_normalize = conf.data.normalize nChunks = np.ceil(len(files_list) / conf.inference.chunk_size).astype(int) + list_of_predictions = [] + list_of_reference_points = [] + ml_super_frame = np.zeros((NY*BinningFactor, NX*BinningFactor), dtype=np.int32) count_frame = np.zeros((NY, NX), dtype=np.int32) subpixel_dist = np.zeros((BinningFactor, BinningFactor), dtype=np.int32) @@ -191,8 +225,30 @@ if __name__ == "__main__": ) predictions, reference_points = run_inference(model, dataLoader, conf) - ml_super_frame_chunk, count_frame_chunk, subpixel_dist_chunk = accumulate_hits(predictions, reference_points, BinningFactor) - ml_super_frame += ml_super_frame_chunk - count_frame += count_frame_chunk - subpixel_dist += subpixel_dist_chunk - save_results(ml_super_frame, count_frame, subpixel_dist, roi, BinningFactor, output_dir) + list_of_predictions.append(predictions) + list_of_reference_points.append(reference_points) + + del dataset, dataLoader + torch.cuda.empty_cache() + + predictions = np.concatenate(list_of_predictions, axis=0) + ref_points = np.concatenate(list_of_reference_points, axis=0) + + del list_of_predictions, list_of_reference_points + + ml_super_frames, count_frame, subpixel_dist = accumulate_hits( + predictions, ref_points, binning_factor=BinningFactor + ) + + print('[Main]: Applying eta interpolation to predictions...') + predictions_eta_interpolated = interpolate_eta_from_points(predictions) + ml_super_frames_eta, count_frame_eta, subpixel_dist_eta = accumulate_hits( + predictions_eta_interpolated, ref_points, binning_factor=BinningFactor + ) + + + save_results(ml_super_frames, ml_super_frames_eta, + count_frame, subpixel_dist, subpixel_dist_eta, + roi=roi, binning_factor=BinningFactor, + output_dir=output_dir + ) \ No newline at end of file