process.py

import argparse
import datetime
import logging
import os
import yaml
import pandas as pd

import generate_masks
import generate_cell_crops
import image_utils
import hpacellseg.cellsegmentator as cellsegmentator


# This is the log configuration. It will log everything to a file AND the console
logging.basicConfig(
    filename="log.txt",
    encoding="utf-8",
    format="%(levelname)s: %(message)s",
    filemode="w",
    level=logging.INFO,
)
console = logging.StreamHandler()
logging.getLogger().addHandler(console)
logger = logging.getLogger("HPACellSegmentatorPortable")

# This is the general configuration variable. We are going to use the special key "log" in the dictionary to use the log in our code
config = {"log": logger}

# If you want to use constants with your script, add them here
config["crop_cells"] = True
config["crop_size"] = 1024
config["crop_bitdepth"] = 8
config["crop_mask"] = True
config["mask_cell"] = True

# If you want to use a configuration file with your script, add it here
with open("config.yaml", "r") as file:
    config_contents = yaml.safe_load(file)
    if config_contents:
        config = config | config_contents

# If you want to use command line parameters with your script, add them here.
# Do NOT set defaults here: the defaults live in the constants/config.yaml layers above.
# We only merge the arguments the user actually passed (value is not None), so the
# priority ends up being constants < config.yaml < command line arguments.
argparser = argparse.ArgumentParser(description="Please input the following parameters")
argparser.add_argument("-c", "--crop_cells", help="if you want to generate the crops of the cells detected in the segmentation", action=argparse.BooleanOptionalAction)
argparser.add_argument("-cs", "--crop_size", help="the cell crop size", type=int)
argparser.add_argument("-cb", "--crop_bitdepth", help="the cell crop bitdepth", type=int)
argparser.add_argument("-cm", "--crop_mask", help="if you want to also generate the crop binary mask from the segmentation", action=argparse.BooleanOptionalAction)
argparser.add_argument("-mc", "--mask_cell", help="if you want additional crops with only the segmented cell area", action=argparse.BooleanOptionalAction)
args = argparser.parse_args()
config = config | {key: value for key, value in vars(args).items() if value is not None}

# Log the start time and the final configuration so you can keep track of what you did
config["log"].info("Start: " + datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S"))
config["log"].info("Parameters used:")
config["log"].info(config)
config["log"].info("----------")

# We load the CellSegmentator model. A failure here aborts the whole run, as there
# is nothing we can process without the model.
try:
    segmentator = cellsegmentator.CellSegmentator(
        "./models/dpn_unet_nuclei_v1.pth",
        "./models/dpn_unet_cell_3ch_v1.pth",
        device="cuda", padding=True, multi_channel_model=True)
except Exception as e:
    config["log"].error("- Could not load the segmentation model: " + str(e))
    raise

# if we want to generate the crops, we are going to keep their information in a CSV file
if config["crop_cells"]:
    df = pd.DataFrame(columns=['id', 'cell', 'x1', 'y1', 'x2', 'y2'])

# We iterate over each set of images to process. Each set is processed inside its own
# try/except so a single failing FOV is logged and skipped instead of aborting the batch.
if os.path.exists("./path_list.csv"):
    with open("./path_list.csv", "r") as path_list:
        for curr_set in path_list:

            if curr_set.strip() != "" and not curr_set.startswith("#"):
                try:
                    curr_set_arr = curr_set.split(",")
                    # We create the output folder
                    os.makedirs(curr_set_arr[4].strip(), exist_ok=True)
                    # We load the images as numpy arrays
                    image_stack = []
                    image_stack.append([image_utils.read_grayscale_image(curr_set_arr[0].strip())])
                    image_stack.append([image_utils.read_grayscale_image(curr_set_arr[1].strip())])
                    image_stack.append([image_utils.read_grayscale_image(curr_set_arr[2].strip())])

                    # We run the model
                    nuclei_mask, cell_mask = generate_masks.create_masks(segmentator, image_stack, curr_set_arr[4].strip(), curr_set_arr[6].strip())
                    # Single cell crops
                    if config["crop_cells"]:
                        os.makedirs(curr_set_arr[5].strip(), exist_ok=True)
                        image_stack.append([image_utils.read_grayscale_image(curr_set_arr[3].strip())])
                        cell_bbox_df = generate_cell_crops.generate_crops(image_stack, cell_mask, nuclei_mask, config["crop_size"], config["crop_bitdepth"], config["crop_mask"], config["mask_cell"], curr_set_arr[5].strip(), curr_set_arr[6].strip())
                        df = pd.concat([df, cell_bbox_df], ignore_index=True)

                    config["log"].info("- Saved results for " + curr_set_arr[6].strip())
                except Exception as e:
                    config["log"].error("- Failed to process '" + curr_set.strip() + "': " + str(e))
                    continue

    # We store the cell crops bboxes and ids for easy localization
    if config["crop_cells"]:
        df.to_csv("crop_info.csv", index=False)

config["log"].info("----------")
config["log"].info("End: " + datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S"))