sum_image.py

import os
import yaml
import numpy as np
import pandas as pd
import tifffile as tif

# same engine as xpdan
from xpdan.glbl import an_glbl
from xpdan.data_reduction import xpd_data_proc
w_dir = os.path.join(an_glbl.home, 'tiff_base')
W_DIR = w_dir

def sum_img(hdr, num_to_sum):
    """function to sum p images from a stack of q images pulled out from
    one hearder, into N images where N= q/p.

    Parameters
    ----------
    hdr : databroker.Header
        header object
    num_to_sum : int
        number of images going to be summed. If images stacks is not
        devisible by designed number of images, the reminding frames
        will be disgarded
    """
    h_start = hdr.get('start', None)
    if not h_start:
        print("WARNING: this header has no start document, "
              "please make sure the scan was completed!")
        return
    expo_time = h_start.get('sp_computed_exposure', None)
    if not expo_time:
        print("WARNING: currently only support scan run with "
              "xpdAcq (wink)")
        return
    total_expo = expo_time*num_to_sum
    print("INFO: exposure time per event = {}s, will be summed "
          "into {}*{} = {:.3f}s".format(expo_time, expo_time,
                                        num_to_sum,
                                        total_expo))
    dark_img, junk = xpd_data_proc.pull_dark(hdr)
    cc = count()
    img = np.zeros_like(dark_img)
    for ev in evs:
        c_num = next(cc) # outer loop
        _img, event_timestamp, ind, is_dark_sub = xpd_data_proc.\
                _dark_sub(ev, dark_img)
        img += _img  # sum image
        if c_num+1 == num_to_sum:  # +1 from python logic
            # save image
            print("INFO: Sum {} images into single frame".format(c_num+1))
            f_name = xpd_data_proc._file_name(ev, event_timestamp, ind)

            root = h_start.get('sample_name', None)
            if root is not None:
                root_dir = os.path.join(W_DIR, 'sum_'+root)
                os.makedirs(root_dir, exist_ok=True)
            else:
                root_dir = W_DIR
            w_name = os.path.join(root_dir,
                                  'sum_{}_{}_{}'.format(num_to_sum,
                                                        total_expo,
                                                        f_name))
            tif.imsave(w_name, img)
            if os.path.isfile(w_name):
                print('image "%s" has been saved at "%s"'
                        % (f_name, root_dir))
            # save run_start
            stem, ext = os.path.splitext(w_name)
            config_name = w_name.replace(ext, '.yml')
            with open(config_name, 'w') as f:
                yaml.dump(h_start, f)  # save all md in start
            # re-init
            img = np.zeros_like(dark_img)
            cc = count() # make new counter