from coseda.logging_utils import log_print
import os
import numpy as np
from concurrent.futures import ProcessPoolExecutor, as_completed
import h5py
from diffractem.peakfinder8_extension import peakfinder_8
import configparser
import datetime
import dask.array as da
from tqdm import tqdm
from coseda.initialize import find_configfiles
from coseda.io import handle_input, parse_config
from coseda.logging_utils import log_start, log_result, shoutout
from coseda.nexus.paths import get_mask_dataset
from coseda.nexus.peaks import ensure_peak_nxdata
from coseda.nexus.process import write_nxprocess_peakfinding
from coseda.peakfinding.maxres import write_maxres_dataset
[docs]
def process_single_frame(h5file_path, frame_index, x0, y0, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res):
with h5py.File(h5file_path, 'r') as file:
image_data = file['entry/data/images'][frame_index, :, :]
shapex = file['entry/data/images'].shape[1]
shapey = file['entry/data/images'].shape[2]
X, Y = np.meshgrid(range(image_data.shape[1]), range(image_data.shape[0]))
R = np.sqrt((X-x0)**2 + (Y-y0)**2).astype(np.float32)
pixmask = None
pm_ds = get_mask_dataset(file)
if pm_ds is not None:
# Only accept 2D masks matching image dimensions
if pm_ds.ndim == 2 and pm_ds.shape == (shapex, shapey):
# Read into memory and cast
pixmask = pm_ds[()]
pixmask = pixmask.astype(np.int8)
# Exclude pixels that are too far or close to the center to be peaks
mask = np.ones_like(image_data, dtype=np.int8)
mask[R > max_res] = 0
mask[R < min_res] = 0
if pixmask is not None:
mask = mask * pixmask
pks = peakfinder_8(500, image_data.astype(np.float32), mask, R, image_data.shape[1], image_data.shape[0], 1, 1, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius)
nPeaks = len(pks[0])
if pks is None or nPeaks == 0:
return []
return list(zip(pks[0], pks[1])) # Return a list of tuples (x, y) for peaks
[docs]
def process_batch(images_batch, x0, y0, num_threads, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res, logfile_path, h5file_path):
#print(f"Starting processing of {len(images_batch)} images in batch.")
batch_results = []
with ProcessPoolExecutor(max_workers=num_threads) as executor:
futures = []
for i, image_data in enumerate(images_batch):
future = executor.submit(process_image, i, image_data, x0, y0, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res)
futures.append(future)
for future in as_completed(futures):
result = future.result()
if result:
batch_results.append(result)
#print(f"Completed processing of {len(images_batch)} images in batch.")
return batch_results
[docs]
def process_image(i, image_data, x0, y0, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res):
nPeaks = 0
X, Y = np.meshgrid(range(image_data.shape[1]), range(image_data.shape[0]))
R = np.sqrt((X-x0)**2 + (Y-y0)**2).astype(np.float32)
mask = np.ones_like(image_data, dtype=np.int8)
mask[R > max_res] = 0
mask[R < min_res] = 0
pks = peakfinder_8(500, image_data.astype(np.float32), mask, R, image_data.shape[1], image_data.shape[0], 1, 1, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius)
nPeaks = len(pks[0])
#print(nPeaks)
if pks is None or len(pks[0]) == 0:
fill = [0] * (500)
return {
'index': i,
'nPeaks': 0,
'peakTotalIntensity': np.array(fill),
'peakXPosRaw': np.array(fill),
'peakYPosRaw': np.array(fill),
}
fill = [0] * (500 - nPeaks)
return {
'index': i,
'nPeaks': nPeaks,
'peakTotalIntensity': np.array(pks[2] + fill),
'peakXPosRaw': np.array(pks[0] + fill),
'peakYPosRaw': np.array(pks[1] + fill),
}
[docs]
def process_file(h5file_path, x0, y0, num_threads, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res, logfile_path, batch_size):
with h5py.File(h5file_path, 'r+') as f:
write_nxprocess_peakfinding(
f,
{
"threshold": threshold,
"min_snr": min_snr,
"min_pix_count": min_pix_count,
"max_pix_count": max_pix_count,
"local_bg_radius": local_bg_radius,
"min_res": min_res,
"max_res": max_res,
"x0": x0,
"y0": y0,
"batch_size": batch_size,
"num_threads": num_threads,
},
program="coseda.peakfinding.findpeaks",
inputs=["/entry/data/images"],
outputs=[
"/entry/data/nPeaks",
"/entry/data/peakTotalIntensity",
"/entry/data/peakXPosRaw",
"/entry/data/peakYPosRaw",
],
)
log_start(logfile_path, f'start peak finding in {os.path.basename(h5file_path)}, threshold = {threshold}, min_snr = {min_snr}, min_pix_count = {min_pix_count}, max_pix_count = {max_pix_count}, local_bg_radius = {local_bg_radius}, min_res = {min_res}, max_res = {max_res}, num_threads = {num_threads}')
total_process_start_time = datetime.datetime.now()
dask_images = da.from_array(f['entry/data/images'], chunks=(batch_size, -1, -1))
dataset_length = len(dask_images)
num_batches = dask_images.shape[0] // batch_size
if dask_images.shape[0] % batch_size != 0:
num_batches += 1
log_start(logfile_path, f'processing in {num_batches} batches of {batch_size} frames')
# Delete datasets if they already exist
for dataset_name in ['nPeaks', 'peakTotalIntensity', 'peakXPosRaw', 'peakYPosRaw', 'index']:
full_name = f'entry/data/{dataset_name}'
if full_name in f:
del f[full_name]
# Create or resize datasets
for dataset_name in ['nPeaks', 'index']:
full_name = f'entry/data/{dataset_name}'
if full_name not in f:
f.create_dataset(full_name, shape=(dask_images.shape[0],), dtype=int)
for dataset_name in ['peakTotalIntensity', 'peakXPosRaw', 'peakYPosRaw']:
full_name = f'entry/data/{dataset_name}'
if full_name not in f:
f.create_dataset(full_name, shape=(dask_images.shape[0], 500), dtype=float) # 2D shape
for batch_idx in range(num_batches):
batch_start_time = datetime.datetime.now()
start_idx = batch_idx * batch_size
end_idx = min(start_idx + batch_size, dask_images.shape[0])
#print(f"processing batch {batch_idx+1}/{num_batches}, images {start_idx} to {end_idx-1}")
log_start(logfile_path,f'start processing batch {batch_idx+1}/{num_batches}, images {start_idx} to {end_idx-1}')
# Use Dask's compute() method to load the batch into memory
images_batch = dask_images[start_idx:end_idx].compute()
# Process each batch
batch_results = process_batch(images_batch, x0, y0, num_threads, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res, logfile_path, h5file_path)
# Estimate the time of completion
batch_end_time = datetime.datetime.now() # Record end time of the current batch
runtime = batch_end_time - total_process_start_time
remaining_batches = num_batches - (batch_idx + 1)
processed_frames = (batch_idx + 1) * batch_size
remaining_frames = dataset_length - processed_frames
estimated_remaining_time = runtime / processed_frames * remaining_frames
etc = datetime.datetime.now() + estimated_remaining_time
log_start(logfile_path,f'finished processing batch {batch_idx+1}/{num_batches}, images {start_idx} to {end_idx-1}')
log_print(f'estimated time of completion: {etc}')
# Write results back to file
for result in batch_results:
# Adjust index for batch offset
idx = result['index'] + start_idx
f['entry/data/nPeaks'][idx] = result['nPeaks']
truncated_peak_intensity = result['peakTotalIntensity'][:500]
f['entry/data/peakTotalIntensity'][idx] = truncated_peak_intensity
truncated_peakXPosRaw = result['peakXPosRaw'][:500]
f['entry/data/peakXPosRaw'][idx] = truncated_peakXPosRaw
truncated_peakYPosRaw = result['peakYPosRaw'][:500]
f['entry/data/peakYPosRaw'][idx] = truncated_peakYPosRaw
f['entry/data/index'][idx] = idx
ensure_peak_nxdata(f)
log_start(logfile_path,f'processing of {os.path.basename(h5file_path)} complete')
# Save per-frame furthest-peak distance in pixels for downstream frame ordering.
write_maxres_dataset(
h5file_path=h5file_path,
center_x=x0,
center_y=y0,
logfile_path=logfile_path,
)
[docs]
def findpeaks8_batch(input_path, batch_size=5000):
configfiles, input_path = handle_input(input_path)
log_print(f'following .ini files were found and will be processed:')
log_print(f'{configfiles}')
counter = 1
for configfile in configfiles:
config, outputfolder, originalfile, logfile, path, outputfolder_path, originalfile_path, logfile_path, framepath, h5file, h5file_path = parse_config(configfile)
#print(f"Working with {os.path.basename(configfile_path)}")
# Check if necessary parameters are defined
for param in ['peakfinding_threshold','peakfinding_min_snr','peakfinding_min_pix_count','peakfinding_max_pix_count','peakfinding_local_bg_radius','peakfinding_min_res','peakfinding_x0','peakfinding_y0','peakfinding_num_threads']:
if not config.has_option('Parameters', param):
with open(f'{logfile_path}', 'a') as file:
file.write(f'{datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]}; Error: {param.split("_", 1)[1]} not defined, peak finding interrupted\n')
raise Exception(f'{param.split("_", 1)[1]} not defined')
# Load parameters
threshold = float(config.get('Parameters', 'peakfinding_threshold'))
min_snr = float(config.get('Parameters', 'peakfinding_min_snr'))
min_pix_count = float(config.get('Parameters', 'peakfinding_min_pix_count'))
max_pix_count = float(config.get('Parameters', 'peakfinding_max_pix_count'))
local_bg_radius = float(config.get('Parameters', 'peakfinding_local_bg_radius'))
min_res = float(config.get('Parameters', 'peakfinding_min_res'))
max_res = float(config.get('Parameters', 'peakfinding_max_res'))
x0 = int(config.get('Parameters', 'peakfinding_x0'))
y0 = int(config.get('Parameters', 'peakfinding_y0'))
num_threads = int(config.get('Parameters', 'peakfinding_num_threads'))
process_file(h5file_path, x0, y0, num_threads, threshold, min_snr, min_pix_count, max_pix_count, local_bg_radius, min_res, max_res, logfile_path, batch_size)
if counter <= len(configfiles):
log_print(f"proceeding to next task (file {counter+1}/{len(configfiles)})")
log_print("")
counter = counter + 1
log_print(f"batch finished ({counter}/{len(configfiles)})")