266 lines
11 KiB
Python
266 lines
11 KiB
Python
import sys
|
|
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
|
|
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):
|
|
output_base = Path(conf.experiment.output_base) / conf.experiment.name / (f'1ph_{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
|
|
|
|
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'[Inference]: mean x = {torch.mean(all_predictions[:, 0]):.4f}, std x = {torch.std(all_predictions[:, 0]):.4f}')
|
|
print(f'[Inference]: 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, number_of_subframes: int = 20):
|
|
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)
|
|
|
|
chunk_size = len(predictions) // number_of_subframes
|
|
|
|
for i in range(number_of_subframes):
|
|
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,
|
|
count_frame, subpixel_dist, subpixel_dist_eta,
|
|
roi: list, binning_factor: int, output_dir: Path):
|
|
x_min, x_max, y_min, y_max = roi
|
|
|
|
# 1. super-resolution frame
|
|
np.save(output_dir / '1Photon_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_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. super-resolution frame with eta interpolation
|
|
np.save(output_dir / '1Photon_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_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 with Eta Interpolation')
|
|
plt.xlabel('X (pixel)')
|
|
plt.ylabel('Y (pixel)')
|
|
plt.savefig(output_dir / '1Photon_ML_superFrame_etaInterpolated.png', dpi=300, bbox_inches='tight')
|
|
plt.close()
|
|
np.save(output_dir / '1Photon_ML_superFrame_etaInterpolated.npy', ml_super_frame_eta)
|
|
|
|
# 3. 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)
|
|
|
|
# 4. 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")
|
|
|
|
# 5. sub-pixel distribution with eta interpolation
|
|
plt.figure(figsize=(8, 8))
|
|
plt.imshow(subpixel_dist_eta, origin='lower',
|
|
extent=[0, binning_factor, 0, binning_factor],
|
|
cmap='viridis')
|
|
plt.colorbar(label='Counts')
|
|
plt.title('Sub-pixel Distribution with Eta Interpolation')
|
|
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_etaInterpolated.png', dpi=300, bbox_inches='tight')
|
|
plt.close()
|
|
np.save(output_dir / '1Photon_subpixel_Distribution_etaInterpolated.npy', subpixel_dist_eta)
|
|
rms_eta, mean_eta = np.std(subpixel_dist_eta), np.mean(subpixel_dist_eta)
|
|
print(f"[Plotting]: Sub-pixel distribution with eta interpolation: RMS/Mean: {rms_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 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))
|
|
|
|
list_of_predictions = []
|
|
list_of_reference_points = []
|
|
|
|
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)
|
|
list_of_predictions.append(_predictions)
|
|
list_of_reference_points.append(_ref_points)
|
|
|
|
del dataset, data_loader
|
|
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
|
|
) |