Upload vista2d version 0.3.1

fd4ffa6 verified over 1 year ago

16.8 kB

# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import gc
import os
import shutil
import time
import warnings
import zipfile

import imageio.v3 as imageio
import numpy as np
from PIL import Image
from pycocotools.coco import COCO


def min_label_precision(label):
    lm = label.max()

    if lm <= 255:
        label = label.astype(np.uint8)
    elif lm <= 65535:
        label = label.astype(np.uint16)
    else:
        label = label.astype(np.uint32)

    return label


def guess_convert_to_uint16(img, margin=30):
    """
    Guess a multiplier that makes all pixels integers.
    The input img (each channel) is already in the range 0..1, they must have been converted from uint16 integers as image / scale,
    where scale was the unknown max intensity.
    We could guess the scale by looking at unique values: 1/np.min(np.diff(np.unique(im)).
    the hypothesis is that it will be more accurate recovery of the original image,
    instead of doing a simple (img*65535).astype(np.uint16)
    """

    for i in range(img.shape[0]):
        im = img[i]

        if im.any():
            start = time.time()
            imsmall = im[::4, ::4]  # subsample
            # imsmall = im

            scale = int(np.round(1 / np.min(np.diff(np.unique(imsmall)))))  # guessing scale
            test = [
                (np.sum((imsmall * k) % 1)) for k in range(scale - margin, scale + margin)
            ]  # finetune, guess a multiplier that makes all pixels integers
            sid = np.argmin(test)  # fine tune scale

            if scale < 16000 or scale > 16400:
                warnings.warn("scale not in expected range")
                print(
                    "guessing scale",
                    scale,
                    test[margin],
                    "fine tuning scale",
                    scale - margin + sid,
                    "dif",
                    test[sid],
                    "time",
                    time.time() - start,
                )

                scale = 16384
            else:
                scale = scale - margin + sid
            # all the recovered scale values seems to be up to 16384,
            # we can stretch to 65535 (for better visualization, most tiff viewers expect that range)
            scale = min(65535, scale * 4)
            img[i] = im * scale

    img = img.astype(np.uint16)
    return img


def concatenate_masks(mask_dir):
    labeled_mask = None
    i = 0
    for filename in sorted(os.listdir(mask_dir)):
        if filename.endswith(".png"):
            mask = imageio.imread(os.path.join(mask_dir, filename))
            if labeled_mask is None:
                labeled_mask = np.zeros(shape=mask.shape, dtype=np.uint16)
            labeled_mask[mask > 0] = i
            i = i + 1

    if i <= 255:
        labeled_mask = labeled_mask.astype(np.uint8)

    return labeled_mask


def get_filenames_exclude_masks(dir1, target_string):
    filenames = []
    # Combine lists of files from both directories
    files = os.listdir(dir1)
    # Filter files that contain the target string but exclude 'masks'
    filenames = [f for f in files if target_string in f and "masks" not in f]

    return filenames


def remove_overlaps(masks, medians, overlap_threshold=0.75):
    """replace overlapping mask pixels with mask id of closest mask
    if mask fully within another mask, remove it
    masks = Nmasks x Ly x Lx
    """
    cellpix = masks.sum(axis=0)
    igood = np.ones(masks.shape[0], "bool")
    for i in masks.sum(axis=(1, 2)).argsort():
        npix = float(masks[i].sum())
        noverlap = float(masks[i][cellpix > 1].sum())
        if noverlap / npix >= overlap_threshold:
            igood[i] = False
            cellpix[masks[i] > 0] -= 1
            # print(cellpix.min())
    print(f"removing {(~igood).sum()} masks")
    masks = masks[igood]
    medians = medians[igood]
    cellpix = masks.sum(axis=0)
    overlaps = np.array(np.nonzero(cellpix > 1.0)).T
    dists = ((overlaps[:, :, np.newaxis] - medians.T) ** 2).sum(axis=1)
    tocell = np.argmin(dists, axis=1)
    masks[:, overlaps[:, 0], overlaps[:, 1]] = 0
    masks[tocell, overlaps[:, 0], overlaps[:, 1]] = 1

    # labels should be 1 to mask.shape[0]
    masks = masks.astype(int) * np.arange(1, masks.shape[0] + 1, 1, int)[:, np.newaxis, np.newaxis]
    masks = masks.sum(axis=0)
    gc.collect()
    return masks


def livecell_process_files(dataset_dir):
    """
    This function takes in the directory of livecell extracted dataset as input and
    extracts labels from the coco format.
    """

    # "A172", "BT474", "Huh7", "MCF7", "SHSY5Y", "SkBr3", "SKOV3"
    # "BV2" is being skipped, runs into memory constraints
    cell_type_list = ["A172", "BT474", "Huh7", "MCF7", "SHSY5Y", "SkBr3", "SKOV3"]
    for each_cell_tp in cell_type_list:
        for split in ["train", "val", "test"]:
            print(f"Working on split: {split}")

            if split == "test":
                img_path = os.path.join(dataset_dir, "images", "livecell_test_images", each_cell_tp)
                msk_path = os.path.join(dataset_dir, "images", "livecell_test_images", each_cell_tp + "_masks")
            else:
                img_path = os.path.join(dataset_dir, "images", "livecell_train_val_images", each_cell_tp)
                msk_path = os.path.join(dataset_dir, "images", "livecell_train_val_images", each_cell_tp + "_masks")
            if not os.path.exists(msk_path):
                os.makedirs(msk_path)

            # annotation path
            path = os.path.join(
                dataset_dir,
                "livecell-dataset.s3.eu-central-1.amazonaws.com",
                "LIVECell_dataset_2021",
                "annotations",
                "LIVECell_single_cells",
                each_cell_tp.lower(),
                split + ".json",
            )
            annotation = COCO(path)
            # Convert COCO format segmentation to binary mask
            images = annotation.loadImgs(annotation.getImgIds())
            height = []
            width = []
            for index, im in enumerate(images):
                print("Status: {}/{}, Process image: {}".format(index, len(images), im["file_name"]))
                if (
                    im["file_name"] == "BV2_Phase_C4_2_03d00h00m_1.tif"
                    or im["file_name"] == "BV2_Phase_C4_2_03d00h00m_3.tif"
                ):
                    print("Skipping the file: BV2_Phase_C4_2_03d00h00m_1.tif, as it is troublesome")
                    continue
                # load image
                img = Image.open(os.path.join(img_path, im["file_name"])).convert("L")
                height.append(img.size[0])
                width.append(img.size[1])

                # load and display instance annotations
                annids = annotation.getAnnIds(imgIds=im["id"], iscrowd=None)
                anns = annotation.loadAnns(annids)

                medians = []
                masks = []
                k = 0
                for ann in anns:
                    # convert segmentation to binary mask
                    mask = annotation.annToMask(ann)
                    masks.append(mask)
                    ypix, xpix = mask.nonzero()
                    medians.append(np.array([ypix.mean().astype(np.float32), xpix.mean().astype(np.float32)]))
                    k += 1

                masks = np.array(masks).astype(np.int8)
                medians = np.array(medians)
                masks = remove_overlaps(masks, medians, overlap_threshold=0.75)
                gc.collect()

                # ## Create new name for the image and also for the mask and save them as .tif format
                # masks_int32 = masks.astype(np.int32)
                # mask_pil = Image.fromarray(masks_int32, 'I')

                t_filename = im["file_name"]
                # cell_type = t_filename.split('_')[0] #? not used
                new_mask_name = t_filename[:-4] + "_masks.tif"
                # mask_pil.save(os.path.join(msk_path, new_mask_name))
                imageio.imwrite(os.path.join(msk_path, new_mask_name), min_label_precision(masks))
                gc.collect()

            print(f"In total {len(images)} images")


def tissuenet_process_files(dataset_dir):
    """
    This function takes in the directory of TissueNet extracted dataset as input and
    creates tiled images into 4 from each image
    """

    for folder in ["train", "val", "test"]:
        if not os.path.exists(os.path.join(dataset_dir, "tissuenet_1.0", folder)):
            os.mkdir(os.path.join(dataset_dir, "tissuenet_1.0", folder))

    for folder in ["train", "val", "test"]:
        print(f"Working on {folder} directory of tissuenet")
        f_name = f"tissuenet_1.0/tissuenet_v1.0_{folder}.npz"
        dat = np.load(os.path.join(dataset_dir, f_name))
        data = dat["X"]
        labels = dat["y"]
        tissues = dat["tissue_list"]
        platforms = dat["platform_list"]
        tlabels = np.unique(tissues)
        plabels = np.unique(platforms)
        tp = 0
        for t in tlabels:
            for p in plabels:
                ix = ((tissues == t) * (platforms == p)).nonzero()[0]
                tp += 1
                if len(ix) > 0:
                    print(f"Working on {t} {p}")

                    for k, i in enumerate(ix):
                        print(f"Status: {k}/{len(ix)} {tp}/{len(tlabels) * len(plabels)} {t} {p}")
                        img = data[i].transpose(2, 0, 1)
                        label = labels[i][:, :, 0]

                        img = guess_convert_to_uint16(img)  # guess inverse scale and convert to uint16
                        label = min_label_precision(label)

                        if folder == "train":
                            img = img.reshape(2, 2, 256, 2, 256).transpose(0, 1, 3, 2, 4).reshape(2, 4, 256, 256)
                            label = label.reshape(2, 256, 2, 256).transpose(0, 2, 1, 3).reshape(4, 256, 256)

                            zero_channel = np.zeros((1, img.shape[1], img.shape[2], img.shape[3]), dtype=img.dtype)

                            # Concatenate the zero channel with the original array along the first dimension
                            new_array = np.concatenate([img, zero_channel], axis=0)
                            # reshaped_array = np.transpose(new_array, (1, 2, 3, 0))
                            for j in range(4):
                                img_name = f"{folder}/{t}_{p}_{k}_{j}.tif"
                                mask_name = f"{folder}/{t}_{p}_{k}_{j}_masks.tif"
                                imageio.imwrite(os.path.join(dataset_dir, "tissuenet_1.0", img_name), new_array[:, j])
                                imageio.imwrite(os.path.join(dataset_dir, "tissuenet_1.0", mask_name), label[j])
                        else:
                            zero_channel = np.zeros((1, img.shape[1], img.shape[2]), dtype=img.dtype)
                            new_array = np.concatenate([img, zero_channel], axis=0)
                            # reshaped_array = np.transpose(new_array, (1, 2, 0))
                            img_name = f"{folder}/{t}_{p}_{k}.tif"
                            mask_name = f"{folder}/{t}_{p}_{k}_masks.tif"
                            imageio.imwrite(os.path.join(dataset_dir, "tissuenet_1.0", img_name), new_array)
                            imageio.imwrite(os.path.join(dataset_dir, "tissuenet_1.0", mask_name), label)


def kaggle_process_files(dataset_dir):
    """
    This function takes in the directory of kaggle nuclei extracted dataset as input and
    creates a json list with 5 folds.
    Please note that there are some hard-coded directory names as per the original dataset.
    The function creates an instance processed dataset and then a 5 fold json file based on
    the instance processed dataset
    """
    data_dir = os.path.join(dataset_dir, "stage1_train")
    saving_path = os.path.join(dataset_dir, "instance_processed_data")
    if not os.path.exists(saving_path):
        os.mkdir(saving_path)

    # Process the images and create instance masks first
    for idx, subdir in enumerate(os.listdir(data_dir)):
        subdir_path = os.path.join(data_dir, subdir)
        if os.path.isdir(subdir_path):
            images_dir = os.path.join(subdir_path, "images")
            masks_dir = os.path.join(subdir_path, "masks")
            if os.path.isdir(images_dir) and os.path.isdir(masks_dir):
                image_file = os.path.join(images_dir, os.listdir(images_dir)[0])
                filename_prefix = f"kg_bowl_{idx}_"

                mask_data = concatenate_masks(masks_dir)

                # ## Apply channel-wise normalization and use only the first three channels
                # image_data = imageio.imread(image_file)
                # normalized_image = normalize_image(image_data[..., :3])
                # imageio.imwrite(os.path.join(saving_path, f"{filename_prefix}img.tiff"), normalized_image)
                shutil.copyfile(image_file, os.path.join(saving_path, f"{filename_prefix}img.png"))
                imageio.imwrite(os.path.join(saving_path, f"{filename_prefix}img_masks.tiff"), mask_data)


def extract_zip(zip_path, extract_to):
    # Ensure the target directory exists
    print(f"Extracting from: {zip_path}")
    print(f"Extracting to: {extract_to}")

    if not os.path.exists(extract_to):
        os.makedirs(extract_to)

    # Extract all contents of the zip file to the specified directory
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(extract_to)


def main():
    parser = argparse.ArgumentParser(description="Script to process the cell imaging datasets")
    parser.add_argument("--dir", type=str, help="Directory of datasets to process it ...", default="/set/the/path")

    args = parser.parse_args()
    data_root_path = os.path.normpath(args.dir)

    dataset_dict = {
        "cellpose_dataset": ["train.zip", "test.zip"],
        "deepbacs_dataset": ["deepbacs.zip"],
        "kaggle_dataset": ["data-science-bowl-2018.zip"],
        "nips_dataset": ["nips_train.zip", "nips_test.zip"],
        "omnipose_dataset": ["datasets.zip"],
        "tissuenet_dataset": ["tissuenet_v1.0.zip"],
        "livecell_dataset": [
            "livecell-dataset.s3.eu-central-1.amazonaws.com/LIVECell_dataset_2021/images_per_celltype.zip"
        ],
    }

    for key, value in dataset_dict.items():
        dataset_path = os.path.join(data_root_path, key)

        for each_zipped in value:
            in_path = os.path.join(dataset_path, each_zipped)
            try:
                if os.path.exists(in_path):
                    print(f"File exists at: {in_path}")
            except Exception:
                print(f"File: {in_path} was not found")
            out_path = os.path.join(dataset_path)
            extract_zip(in_path, out_path)

    print("If we reached here, that means all zip files got extracted ... Working on pre-processing")

    # Looping over all datasets again, Cellpose & Deepbacs have a similar directory structure
    for key, _value in dataset_dict.items():
        if key == "kaggle_dataset":
            print("Needs additional extraction")
            train_zip_path = os.path.join(data_root_path, key, "stage1_train.zip")
            zip_out_path = os.path.join(data_root_path, key, "stage1_train")
            extract_zip(train_zip_path, zip_out_path)
            print("Processing Kaggle Dataset ...")
            dataset_path = os.path.join(data_root_path, key)
            kaggle_process_files(dataset_dir=dataset_path)

        elif key == "livecell_dataset":
            print("Processing LiveCell Dataset ...")
            print(
                "Fyi, this processing might take upto an hour, coffee break might be more fruitful in the meanwhile ..."
            )
            dataset_path = os.path.join(data_root_path, key)
            livecell_process_files(dataset_dir=dataset_path)

        elif key == "tissuenet_dataset":
            print("Processing TissueNet Dataset ...")
            dataset_path = os.path.join(data_root_path, key)
            tissuenet_process_files(dataset_dir=dataset_path)

    return None


if __name__ == "__main__":
    main()

	# Copyright (c) MONAI Consortium
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	# http://www.apache.org/licenses/LICENSE-2.0
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	import argparse
	import gc
	import os
	import shutil
	import time
	import warnings
	import zipfile

	import imageio.v3 as imageio
	import numpy as np
	from PIL import Image
	from pycocotools.coco import COCO


	def min_label_precision(label):
	lm = label.max()

	if lm <= 255:
	label = label.astype(np.uint8)
	elif lm <= 65535:
	label = label.astype(np.uint16)
	else:
	label = label.astype(np.uint32)

	return label


	def guess_convert_to_uint16(img, margin=30):
	"""
	Guess a multiplier that makes all pixels integers.
	The input img (each channel) is already in the range 0..1, they must have been converted from uint16 integers as image / scale,
	where scale was the unknown max intensity.
	We could guess the scale by looking at unique values: 1/np.min(np.diff(np.unique(im)).
	the hypothesis is that it will be more accurate recovery of the original image,
	instead of doing a simple (img*65535).astype(np.uint16)
	"""

	for i in range(img.shape[0]):
	im = img[i]

	if im.any():
	start = time.time()
	imsmall = im[::4, ::4] # subsample
	# imsmall = im

	scale = int(np.round(1 / np.min(np.diff(np.unique(imsmall))))) # guessing scale
	test = [
	(np.sum((imsmall * k) % 1)) for k in range(scale - margin, scale + margin)
	] # finetune, guess a multiplier that makes all pixels integers
	sid = np.argmin(test) # fine tune scale

	if scale < 16000 or scale > 16400:
	warnings.warn("scale not in expected range")
	print(
	"guessing scale",
	scale,
	test[margin],
	"fine tuning scale",
	scale - margin + sid,
	"dif",
	test[sid],
	"time",
	time.time() - start,
	)

	scale = 16384
	else:
	scale = scale - margin + sid
	# all the recovered scale values seems to be up to 16384,
	# we can stretch to 65535 (for better visualization, most tiff viewers expect that range)
	scale = min(65535, scale * 4)
	img[i] = im * scale

	img = img.astype(np.uint16)
	return img


	def concatenate_masks(mask_dir):
	labeled_mask = None
	i = 0
	for filename in sorted(os.listdir(mask_dir)):
	if filename.endswith(".png"):
	mask = imageio.imread(os.path.join(mask_dir, filename))
	if labeled_mask is None:
	labeled_mask = np.zeros(shape=mask.shape, dtype=np.uint16)
	labeled_mask[mask > 0] = i
	i = i + 1

	if i <= 255:
	labeled_mask = labeled_mask.astype(np.uint8)

	return labeled_mask


	def get_filenames_exclude_masks(dir1, target_string):
	filenames = []
	# Combine lists of files from both directories
	files = os.listdir(dir1)
	# Filter files that contain the target string but exclude 'masks'
	filenames = [f for f in files if target_string in f and "masks" not in f]

	return filenames


	def remove_overlaps(masks, medians, overlap_threshold=0.75):
	"""replace overlapping mask pixels with mask id of closest mask
	if mask fully within another mask, remove it
	masks = Nmasks x Ly x Lx
	"""
	cellpix = masks.sum(axis=0)
	igood = np.ones(masks.shape[0], "bool")
	for i in masks.sum(axis=(1, 2)).argsort():
	npix = float(masks[i].sum())
	noverlap = float(masks[i][cellpix > 1].sum())
	if noverlap / npix >= overlap_threshold:
	igood[i] = False
	cellpix[masks[i] > 0] -= 1
	# print(cellpix.min())
	print(f"removing {(~igood).sum()} masks")
	masks = masks[igood]
	medians = medians[igood]
	cellpix = masks.sum(axis=0)
	overlaps = np.array(np.nonzero(cellpix > 1.0)).T
	dists = ((overlaps[:, :, np.newaxis] - medians.T) ** 2).sum(axis=1)
	tocell = np.argmin(dists, axis=1)
	masks[:, overlaps[:, 0], overlaps[:, 1]] = 0
	masks[tocell, overlaps[:, 0], overlaps[:, 1]] = 1

	# labels should be 1 to mask.shape[0]
	masks = masks.astype(int) * np.arange(1, masks.shape[0] + 1, 1, int)[:, np.newaxis, np.newaxis]
	masks = masks.sum(axis=0)
	gc.collect()
	return masks


	def livecell_process_files(dataset_dir):
	"""
	This function takes in the directory of livecell extracted dataset as input and
	extracts labels from the coco format.
	"""

	# "A172", "BT474", "Huh7", "MCF7", "SHSY5Y", "SkBr3", "SKOV3"
	# "BV2" is being skipped, runs into memory constraints
	cell_type_list = ["A172", "BT474", "Huh7", "MCF7", "SHSY5Y", "SkBr3", "SKOV3"]
	for each_cell_tp in cell_type_list:
	for split in ["train", "val", "test"]:
	print(f"Working on split: {split}")

	if split == "test":
	img_path = os.path.join(dataset_dir, "images", "livecell_test_images", each_cell_tp)
	msk_path = os.path.join(dataset_dir, "images", "livecell_test_images", each_cell_tp + "_masks")
	else:
	img_path = os.path.join(dataset_dir, "images", "livecell_train_val_images", each_cell_tp)
	msk_path = os.path.join(dataset_dir, "images", "livecell_train_val_images", each_cell_tp + "_masks")
	if not os.path.exists(msk_path):
	os.makedirs(msk_path)

	# annotation path
	path = os.path.join(
	dataset_dir,
	"livecell-dataset.s3.eu-central-1.amazonaws.com",
	"LIVECell_dataset_2021",
	"annotations",
	"LIVECell_single_cells",
	each_cell_tp.lower(),
	split + ".json",
	)
	annotation = COCO(path)
	# Convert COCO format segmentation to binary mask
	images = annotation.loadImgs(annotation.getImgIds())
	height = []
	width = []
	for index, im in enumerate(images):
	print("Status: {}/{}, Process image: {}".format(index, len(images), im["file_name"]))
	if (
	im["file_name"] == "BV2_Phase_C4_2_03d00h00m_1.tif"
	or im["file_name"] == "BV2_Phase_C4_2_03d00h00m_3.tif"
	):
	print("Skipping the file: BV2_Phase_C4_2_03d00h00m_1.tif, as it is troublesome")
	continue
	# load image
	img = Image.open(os.path.join(img_path, im["file_name"])).convert("L")
	height.append(img.size[0])
	width.append(img.size[1])

	# load and display instance annotations
	annids = annotation.getAnnIds(imgIds=im["id"], iscrowd=None)
	anns = annotation.loadAnns(annids)

	medians = []
	masks = []
	k = 0
	for ann in anns:
	# convert segmentation to binary mask
	mask = annotation.annToMask(ann)
	masks.append(mask)
	ypix, xpix = mask.nonzero()
	medians.append(np.array([ypix.mean().astype(np.float32), xpix.mean().astype(np.float32)]))
	k += 1

	masks = np.array(masks).astype(np.int8)
	medians = np.array(medians)
	masks = remove_overlaps(masks, medians, overlap_threshold=0.75)
	gc.collect()

	# ## Create new name for the image and also for the mask and save them as .tif format
	# masks_int32 = masks.astype(np.int32)
	# mask_pil = Image.fromarray(masks_int32, 'I')

	t_filename = im["file_name"]
	# cell_type = t_filename.split('_')[0] #? not used
	new_mask_name = t_filename[:-4] + "_masks.tif"
	# mask_pil.save(os.path.join(msk_path, new_mask_name))
	imageio.imwrite(os.path.join(msk_path, new_mask_name), min_label_precision(masks))
	gc.collect()

	print(f"In total {len(images)} images")


	def tissuenet_process_files(dataset_dir):
	"""
	This function takes in the directory of TissueNet extracted dataset as input and
	creates tiled images into 4 from each image
	"""

	for folder in ["train", "val", "test"]:
	if not os.path.exists(os.path.join(dataset_dir, "tissuenet_1.0", folder)):
	os.mkdir(os.path.join(dataset_dir, "tissuenet_1.0", folder))

	for folder in ["train", "val", "test"]:
	print(f"Working on {folder} directory of tissuenet")
	f_name = f"tissuenet_1.0/tissuenet_v1.0_{folder}.npz"
	dat = np.load(os.path.join(dataset_dir, f_name))
	data = dat["X"]
	labels = dat["y"]
	tissues = dat["tissue_list"]
	platforms = dat["platform_list"]
	tlabels = np.unique(tissues)
	plabels = np.unique(platforms)
	tp = 0
	for t in tlabels:
	for p in plabels:
	ix = ((tissues == t) * (platforms == p)).nonzero()[0]
	tp += 1
	if len(ix) > 0:
	print(f"Working on {t} {p}")

	for k, i in enumerate(ix):
	print(f"Status: {k}/{len(ix)} {tp}/{len(tlabels) * len(plabels)} {t} {p}")
	img = data[i].transpose(2, 0, 1)
	label = labels[i][:, :, 0]

	img = guess_convert_to_uint16(img) # guess inverse scale and convert to uint16
	label = min_label_precision(label)

	if folder == "train":
	img = img.reshape(2, 2, 256, 2, 256).transpose(0, 1, 3, 2, 4).reshape(2, 4, 256, 256)
	label = label.reshape(2, 256, 2, 256).transpose(0, 2, 1, 3).reshape(4, 256, 256)

	zero_channel = np.zeros((1, img.shape[1], img.shape[2], img.shape[3]), dtype=img.dtype)

	# Concatenate the zero channel with the original array along the first dimension
	new_array = np.concatenate([img, zero_channel], axis=0)
	# reshaped_array = np.transpose(new_array, (1, 2, 3, 0))
	for j in range(4):
	img_name = f"{folder}/{t}_{p}_{k}_{j}.tif"
	mask_name = f"{folder}/{t}_{p}_{k}_{j}_masks.tif"
	imageio.imwrite(os.path.join(dataset_dir, "tissuenet_1.0", img_name), new_array[:, j])
	imageio.imwrite(os.path.join(dataset_dir, "tissuenet_1.0", mask_name), label[j])
	else:
	zero_channel = np.zeros((1, img.shape[1], img.shape[2]), dtype=img.dtype)
	new_array = np.concatenate([img, zero_channel], axis=0)
	# reshaped_array = np.transpose(new_array, (1, 2, 0))
	img_name = f"{folder}/{t}_{p}_{k}.tif"
	mask_name = f"{folder}/{t}_{p}_{k}_masks.tif"
	imageio.imwrite(os.path.join(dataset_dir, "tissuenet_1.0", img_name), new_array)
	imageio.imwrite(os.path.join(dataset_dir, "tissuenet_1.0", mask_name), label)


	def kaggle_process_files(dataset_dir):
	"""
	This function takes in the directory of kaggle nuclei extracted dataset as input and
	creates a json list with 5 folds.
	Please note that there are some hard-coded directory names as per the original dataset.
	The function creates an instance processed dataset and then a 5 fold json file based on
	the instance processed dataset
	"""
	data_dir = os.path.join(dataset_dir, "stage1_train")
	saving_path = os.path.join(dataset_dir, "instance_processed_data")
	if not os.path.exists(saving_path):
	os.mkdir(saving_path)

	# Process the images and create instance masks first
	for idx, subdir in enumerate(os.listdir(data_dir)):
	subdir_path = os.path.join(data_dir, subdir)
	if os.path.isdir(subdir_path):
	images_dir = os.path.join(subdir_path, "images")
	masks_dir = os.path.join(subdir_path, "masks")
	if os.path.isdir(images_dir) and os.path.isdir(masks_dir):
	image_file = os.path.join(images_dir, os.listdir(images_dir)[0])
	filename_prefix = f"kg_bowl_{idx}_"

	mask_data = concatenate_masks(masks_dir)

	# ## Apply channel-wise normalization and use only the first three channels
	# image_data = imageio.imread(image_file)
	# normalized_image = normalize_image(image_data[..., :3])
	# imageio.imwrite(os.path.join(saving_path, f"{filename_prefix}img.tiff"), normalized_image)
	shutil.copyfile(image_file, os.path.join(saving_path, f"{filename_prefix}img.png"))
	imageio.imwrite(os.path.join(saving_path, f"{filename_prefix}img_masks.tiff"), mask_data)


	def extract_zip(zip_path, extract_to):
	# Ensure the target directory exists
	print(f"Extracting from: {zip_path}")
	print(f"Extracting to: {extract_to}")

	if not os.path.exists(extract_to):
	os.makedirs(extract_to)

	# Extract all contents of the zip file to the specified directory
	with zipfile.ZipFile(zip_path, "r") as zip_ref:
	zip_ref.extractall(extract_to)


	def main():
	parser = argparse.ArgumentParser(description="Script to process the cell imaging datasets")
	parser.add_argument("--dir", type=str, help="Directory of datasets to process it ...", default="/set/the/path")

	args = parser.parse_args()
	data_root_path = os.path.normpath(args.dir)

	dataset_dict = {
	"cellpose_dataset": ["train.zip", "test.zip"],
	"deepbacs_dataset": ["deepbacs.zip"],
	"kaggle_dataset": ["data-science-bowl-2018.zip"],
	"nips_dataset": ["nips_train.zip", "nips_test.zip"],
	"omnipose_dataset": ["datasets.zip"],
	"tissuenet_dataset": ["tissuenet_v1.0.zip"],
	"livecell_dataset": [
	"livecell-dataset.s3.eu-central-1.amazonaws.com/LIVECell_dataset_2021/images_per_celltype.zip"
	],
	}

	for key, value in dataset_dict.items():
	dataset_path = os.path.join(data_root_path, key)

	for each_zipped in value:
	in_path = os.path.join(dataset_path, each_zipped)
	try:
	if os.path.exists(in_path):
	print(f"File exists at: {in_path}")
	except Exception:
	print(f"File: {in_path} was not found")
	out_path = os.path.join(dataset_path)
	extract_zip(in_path, out_path)

	print("If we reached here, that means all zip files got extracted ... Working on pre-processing")

	# Looping over all datasets again, Cellpose & Deepbacs have a similar directory structure
	for key, _value in dataset_dict.items():
	if key == "kaggle_dataset":
	print("Needs additional extraction")
	train_zip_path = os.path.join(data_root_path, key, "stage1_train.zip")
	zip_out_path = os.path.join(data_root_path, key, "stage1_train")
	extract_zip(train_zip_path, zip_out_path)
	print("Processing Kaggle Dataset ...")
	dataset_path = os.path.join(data_root_path, key)
	kaggle_process_files(dataset_dir=dataset_path)

	elif key == "livecell_dataset":
	print("Processing LiveCell Dataset ...")
	print(
	"Fyi, this processing might take upto an hour, coffee break might be more fruitful in the meanwhile ..."
	)
	dataset_path = os.path.join(data_root_path, key)
	livecell_process_files(dataset_dir=dataset_path)

	elif key == "tissuenet_dataset":
	print("Processing TissueNet Dataset ...")
	dataset_path = os.path.join(data_root_path, key)
	tissuenet_process_files(dataset_dir=dataset_path)

	return None


	if __name__ == "__main__":
	main()