Upload vista2d version 0.3.1

fd4ffa6 verified over 1 year ago

11.7 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 json
import os

import cv2
import fastremap
import numpy as np
import PIL
import tifffile
import torch
import torch.nn.functional as F
from cellpose.dynamics import compute_masks, masks_to_flows
from cellpose.metrics import _intersection_over_union, _true_positive
from monai.apps import get_logger
from monai.data import MetaTensor
from monai.transforms import MapTransform
from monai.utils import ImageMetaKey, convert_to_dst_type

logger = get_logger("VistaCell")


class LoadTiffd(MapTransform):
    def __call__(self, data):
        d = dict(data)
        for key in self.key_iterator(d):
            filename = d[key]

            extension = os.path.splitext(filename)[1][1:]
            image_size = None

            if extension in ["tif", "tiff"]:
                img_array = tifffile.imread(filename)  # use tifffile for tif images
                image_size = img_array.shape
                if len(img_array.shape) == 3 and img_array.shape[-1] <= 3:
                    img_array = np.transpose(img_array, (2, 0, 1))  # channels first without transpose
            else:
                img_array = np.array(PIL.Image.open(filename))  # PIL for all other images (png, jpeg)
                image_size = img_array.shape
                if len(img_array.shape) == 3:
                    img_array = np.transpose(img_array, (2, 0, 1))  # channels first

            if len(img_array.shape) not in [2, 3]:
                raise ValueError(
                    "Unsupported image dimensions, filename " + str(filename) + " shape " + str(img_array.shape)
                )

            if len(img_array.shape) == 2:
                img_array = img_array[np.newaxis]  # add channels_first if no channel

            if key == "label":
                if img_array.shape[0] > 1:
                    print(
                        f"Strange case, label with several channels {filename} shape {img_array.shape}, keeping only first"
                    )
                    img_array = img_array[[0]]

            elif key == "image":
                if img_array.shape[0] == 1:
                    img_array = np.repeat(img_array, repeats=3, axis=0)  # if grayscale, repeat as 3 channels
                elif img_array.shape[0] == 2:
                    print(
                        f"Strange case, image with 2 channels {filename} shape {img_array.shape}, appending first channel to make 3"
                    )
                    img_array = np.stack(
                        (img_array[0], img_array[1], img_array[0]), axis=0
                    )  # this should not happen, we got 2 channel input image
                elif img_array.shape[0] > 3:
                    print(f"Strange case, image with >3 channels,  {filename} shape {img_array.shape}, keeping first 3")
                    img_array = img_array[:3]

            meta_data = {ImageMetaKey.FILENAME_OR_OBJ: filename, ImageMetaKey.SPATIAL_SHAPE: image_size}
            d[key] = MetaTensor.ensure_torch_and_prune_meta(img_array, meta_data)

        return d


class SaveTiffd(MapTransform):
    def __init__(self, output_dir, data_root_dir="/", nested_folder=False, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self.output_dir = output_dir
        self.data_root_dir = data_root_dir
        self.nested_folder = nested_folder

    def set_data_root_dir(self, data_root_dir):
        self.data_root_dir = data_root_dir

    def __call__(self, data):
        d = dict(data)
        os.makedirs(self.output_dir, exist_ok=True)

        for key in self.key_iterator(d):
            seg = d[key]
            filename = seg.meta[ImageMetaKey.FILENAME_OR_OBJ]

            basename = os.path.splitext(os.path.basename(filename))[0]

            if self.nested_folder:
                reldir = os.path.relpath(os.path.dirname(filename), self.data_root_dir)
                outdir = os.path.join(self.output_dir, reldir)
                os.makedirs(outdir, exist_ok=True)
            else:
                outdir = self.output_dir

            outname = os.path.join(outdir, basename + ".tif")

            label = seg.cpu().numpy()
            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)

            tifffile.imwrite(outname, label)

            print(f"Saving {outname} shape {label.shape} max {label.max()} dtype {label.dtype}")

        return d


class LabelsToFlows(MapTransform):
    # based on dynamics labels_to_flows()
    # created a 3 channel output (foreground, flowx, flowy) and saves under flow (new) key

    def __init__(self, flow_key, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.flow_key = flow_key

    def __call__(self, data):
        d = dict(data)
        for key in self.key_iterator(d):
            label = d[key].int().numpy()

            label = fastremap.renumber(label, in_place=True)[0]
            veci = masks_to_flows(label[0], device=None)

            flows = np.concatenate((label > 0.5, veci), axis=0).astype(np.float32)
            flows = convert_to_dst_type(flows, d[key], dtype=torch.float, device=d[key].device)[0]
            d[self.flow_key] = flows
            # meta_data = {ImageMetaKey.FILENAME_OR_OBJ : filename}
            # d[key] = MetaTensor.ensure_torch_and_prune_meta(img_array, meta_data)
        return d


class LogitsToLabels:
    def __call__(self, logits, filename=None):
        device = logits.device
        logits = logits.float().cpu().numpy()
        dp = logits[1:]  # vectors
        cellprob = logits[0]  # foreground prob (logit)

        try:
            pred_mask, p = compute_masks(
                dp, cellprob, niter=200, cellprob_threshold=0.4, flow_threshold=0.4, interp=True, device=device
            )
        except RuntimeError as e:
            logger.warning(f"compute_masks failed on GPU retrying on CPU {logits.shape} file {filename} {e}")
            pred_mask, p = compute_masks(
                dp, cellprob, niter=200, cellprob_threshold=0.4, flow_threshold=0.4, interp=True, device=None
            )

        return pred_mask, p


class LogitsToLabelsd(MapTransform):
    def __call__(self, data):
        d = dict(data)
        f = LogitsToLabels()
        for key in self.key_iterator(d):
            pred_mask, p = f(d[key])
            d[key] = pred_mask
            d[f"{key}_centroids"] = p
        return d


class SaveTiffExd(MapTransform):
    def __init__(self, output_dir, output_ext=".png", output_postfix="seg", image_key="image", *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self.output_dir = output_dir
        self.output_ext = output_ext
        self.output_postfix = output_postfix
        self.image_key = image_key

    def to_polygons(self, contours):
        polygons = []
        for contour in contours:
            if len(contour) < 3:
                continue
            polygons.append(np.squeeze(contour).astype(int).tolist())
        return polygons

    def __call__(self, data):
        d = dict(data)

        output_dir = d.get("output_dir", self.output_dir)
        output_ext = d.get("output_ext", self.output_ext)
        overlayed_masks = d.get("overlayed_masks", False)
        output_contours = d.get("output_contours", False)

        os.makedirs(self.output_dir, exist_ok=True)

        img = d.get(self.image_key, None)
        filename = img.meta.get(ImageMetaKey.FILENAME_OR_OBJ) if img is not None else None
        image_size = img.meta.get(ImageMetaKey.SPATIAL_SHAPE) if img is not None else None
        basename = os.path.splitext(os.path.basename(filename))[0] if filename else "mask"
        logger.info(f"File: {filename}; Base: {basename}")

        for key in self.key_iterator(d):
            label = d[key]
            output_filename = f"{basename}{'_' + self.output_postfix if self.output_postfix else ''}{output_ext}"
            output_filepath = os.path.join(output_dir, output_filename)
            lm = label.max()
            logger.info(f"Mask Shape: {label.shape}; Instances: {lm}")

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

            tifffile.imwrite(output_filepath, label)
            logger.info(f"Saving {output_filepath}")

            polygons = []
            if overlayed_masks:
                logger.info(f"Overlay Masks: Reading original Image: {filename}")
                image = cv2.imread(filename)
                mask = cv2.imread(output_filepath, 0)

                for i in range(1, np.max(mask)):
                    m = np.zeros_like(mask)
                    m[mask == i] = 1
                    color = np.random.choice(range(256), size=3).tolist()
                    contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
                    polygons.extend(self.to_polygons(contours))
                    cv2.drawContours(image, contours, -1, color, 1)
                cv2.imwrite(output_filepath, image)
                logger.info(f"Overlay Masks: Saving {output_filepath}")
            else:
                label = cv2.convertScaleAbs(label, alpha=255.0 / label.max())
                contours, _ = cv2.findContours(label, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
                polygons.extend(self.to_polygons(contours))

            meta_json = {"image_size": image_size, "contours": len(polygons)}
            with open(os.path.join(output_dir, "meta.json"), "w") as fp:
                json.dump(meta_json, fp, indent=2)

            if output_contours:
                logger.info(f"Total Polygons: {len(polygons)}")
                with open(os.path.join(output_dir, "contours.json"), "w") as fp:
                    json.dump({"count": len(polygons), "contours": polygons}, fp, indent=2)

        return d


# Loss (adopted from Cellpose)
class CellLoss:
    def __call__(self, y_pred, y):
        loss = 0.5 * F.mse_loss(y_pred[:, 1:], 5 * y[:, 1:]) + F.binary_cross_entropy_with_logits(
            y_pred[:, [0]], y[:, [0]]
        )
        return loss


# Accuracy (adopted from Cellpose)
class CellAcc:
    def __call__(self, mask_pred, mask_true):
        if isinstance(mask_true, torch.Tensor):
            mask_true = mask_true.cpu().numpy()

        if isinstance(mask_pred, torch.Tensor):
            mask_pred = mask_pred.cpu().numpy()

        # print("CellAcc mask_true", mask_true.shape, 'max', np.max(mask_true), ",
        #       "'mask_pred', mask_pred.shape,  'max', np.max(mask_pred) )

        iou = _intersection_over_union(mask_true, mask_pred)[1:, 1:]
        tp = _true_positive(iou, th=0.5)

        fp = np.max(mask_pred) - tp
        fn = np.max(mask_true) - tp
        ap = tp / (tp + fp + fn)

        # print("CellAcc ap", ap, 'tp', tp, 'fp', fp,  'fn', fn)
        return ap

	# 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 json
	import os

	import cv2
	import fastremap
	import numpy as np
	import PIL
	import tifffile
	import torch
	import torch.nn.functional as F
	from cellpose.dynamics import compute_masks, masks_to_flows
	from cellpose.metrics import _intersection_over_union, _true_positive
	from monai.apps import get_logger
	from monai.data import MetaTensor
	from monai.transforms import MapTransform
	from monai.utils import ImageMetaKey, convert_to_dst_type

	logger = get_logger("VistaCell")


	class LoadTiffd(MapTransform):
	def __call__(self, data):
	d = dict(data)
	for key in self.key_iterator(d):
	filename = d[key]

	extension = os.path.splitext(filename)[1][1:]
	image_size = None

	if extension in ["tif", "tiff"]:
	img_array = tifffile.imread(filename) # use tifffile for tif images
	image_size = img_array.shape
	if len(img_array.shape) == 3 and img_array.shape[-1] <= 3:
	img_array = np.transpose(img_array, (2, 0, 1)) # channels first without transpose
	else:
	img_array = np.array(PIL.Image.open(filename)) # PIL for all other images (png, jpeg)
	image_size = img_array.shape
	if len(img_array.shape) == 3:
	img_array = np.transpose(img_array, (2, 0, 1)) # channels first

	if len(img_array.shape) not in [2, 3]:
	raise ValueError(
	"Unsupported image dimensions, filename " + str(filename) + " shape " + str(img_array.shape)
	)

	if len(img_array.shape) == 2:
	img_array = img_array[np.newaxis] # add channels_first if no channel

	if key == "label":
	if img_array.shape[0] > 1:
	print(
	f"Strange case, label with several channels {filename} shape {img_array.shape}, keeping only first"
	)
	img_array = img_array[[0]]

	elif key == "image":
	if img_array.shape[0] == 1:
	img_array = np.repeat(img_array, repeats=3, axis=0) # if grayscale, repeat as 3 channels
	elif img_array.shape[0] == 2:
	print(
	f"Strange case, image with 2 channels {filename} shape {img_array.shape}, appending first channel to make 3"
	)
	img_array = np.stack(
	(img_array[0], img_array[1], img_array[0]), axis=0
	) # this should not happen, we got 2 channel input image
	elif img_array.shape[0] > 3:
	print(f"Strange case, image with >3 channels, {filename} shape {img_array.shape}, keeping first 3")
	img_array = img_array[:3]

	meta_data = {ImageMetaKey.FILENAME_OR_OBJ: filename, ImageMetaKey.SPATIAL_SHAPE: image_size}
	d[key] = MetaTensor.ensure_torch_and_prune_meta(img_array, meta_data)

	return d


	class SaveTiffd(MapTransform):
	def __init__(self, output_dir, data_root_dir="/", nested_folder=False, args, *kwargs) -> None:
	super().__init__(args, *kwargs)

	self.output_dir = output_dir
	self.data_root_dir = data_root_dir
	self.nested_folder = nested_folder

	def set_data_root_dir(self, data_root_dir):
	self.data_root_dir = data_root_dir

	def __call__(self, data):
	d = dict(data)
	os.makedirs(self.output_dir, exist_ok=True)

	for key in self.key_iterator(d):
	seg = d[key]
	filename = seg.meta[ImageMetaKey.FILENAME_OR_OBJ]

	basename = os.path.splitext(os.path.basename(filename))[0]

	if self.nested_folder:
	reldir = os.path.relpath(os.path.dirname(filename), self.data_root_dir)
	outdir = os.path.join(self.output_dir, reldir)
	os.makedirs(outdir, exist_ok=True)
	else:
	outdir = self.output_dir

	outname = os.path.join(outdir, basename + ".tif")

	label = seg.cpu().numpy()
	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)

	tifffile.imwrite(outname, label)

	print(f"Saving {outname} shape {label.shape} max {label.max()} dtype {label.dtype}")

	return d


	class LabelsToFlows(MapTransform):
	# based on dynamics labels_to_flows()
	# created a 3 channel output (foreground, flowx, flowy) and saves under flow (new) key

	def __init__(self, flow_key, args, *kwargs) -> None:
	super().__init__(args, *kwargs)
	self.flow_key = flow_key

	def __call__(self, data):
	d = dict(data)
	for key in self.key_iterator(d):
	label = d[key].int().numpy()

	label = fastremap.renumber(label, in_place=True)[0]
	veci = masks_to_flows(label[0], device=None)

	flows = np.concatenate((label > 0.5, veci), axis=0).astype(np.float32)
	flows = convert_to_dst_type(flows, d[key], dtype=torch.float, device=d[key].device)[0]
	d[self.flow_key] = flows
	# meta_data = {ImageMetaKey.FILENAME_OR_OBJ : filename}
	# d[key] = MetaTensor.ensure_torch_and_prune_meta(img_array, meta_data)
	return d


	class LogitsToLabels:
	def __call__(self, logits, filename=None):
	device = logits.device
	logits = logits.float().cpu().numpy()
	dp = logits[1:] # vectors
	cellprob = logits[0] # foreground prob (logit)

	try:
	pred_mask, p = compute_masks(
	dp, cellprob, niter=200, cellprob_threshold=0.4, flow_threshold=0.4, interp=True, device=device
	)
	except RuntimeError as e:
	logger.warning(f"compute_masks failed on GPU retrying on CPU {logits.shape} file {filename} {e}")
	pred_mask, p = compute_masks(
	dp, cellprob, niter=200, cellprob_threshold=0.4, flow_threshold=0.4, interp=True, device=None
	)

	return pred_mask, p


	class LogitsToLabelsd(MapTransform):
	def __call__(self, data):
	d = dict(data)
	f = LogitsToLabels()
	for key in self.key_iterator(d):
	pred_mask, p = f(d[key])
	d[key] = pred_mask
	d[f"{key}_centroids"] = p
	return d


	class SaveTiffExd(MapTransform):
	def __init__(self, output_dir, output_ext=".png", output_postfix="seg", image_key="image", args, *kwargs) -> None:
	super().__init__(args, *kwargs)

	self.output_dir = output_dir
	self.output_ext = output_ext
	self.output_postfix = output_postfix
	self.image_key = image_key

	def to_polygons(self, contours):
	polygons = []
	for contour in contours:
	if len(contour) < 3:
	continue
	polygons.append(np.squeeze(contour).astype(int).tolist())
	return polygons

	def __call__(self, data):
	d = dict(data)

	output_dir = d.get("output_dir", self.output_dir)
	output_ext = d.get("output_ext", self.output_ext)
	overlayed_masks = d.get("overlayed_masks", False)
	output_contours = d.get("output_contours", False)

	os.makedirs(self.output_dir, exist_ok=True)

	img = d.get(self.image_key, None)
	filename = img.meta.get(ImageMetaKey.FILENAME_OR_OBJ) if img is not None else None
	image_size = img.meta.get(ImageMetaKey.SPATIAL_SHAPE) if img is not None else None
	basename = os.path.splitext(os.path.basename(filename))[0] if filename else "mask"
	logger.info(f"File: {filename}; Base: {basename}")

	for key in self.key_iterator(d):
	label = d[key]
	output_filename = f"{basename}{'_' + self.output_postfix if self.output_postfix else ''}{output_ext}"
	output_filepath = os.path.join(output_dir, output_filename)
	lm = label.max()
	logger.info(f"Mask Shape: {label.shape}; Instances: {lm}")

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

	tifffile.imwrite(output_filepath, label)
	logger.info(f"Saving {output_filepath}")

	polygons = []
	if overlayed_masks:
	logger.info(f"Overlay Masks: Reading original Image: {filename}")
	image = cv2.imread(filename)
	mask = cv2.imread(output_filepath, 0)

	for i in range(1, np.max(mask)):
	m = np.zeros_like(mask)
	m[mask == i] = 1
	color = np.random.choice(range(256), size=3).tolist()
	contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
	polygons.extend(self.to_polygons(contours))
	cv2.drawContours(image, contours, -1, color, 1)
	cv2.imwrite(output_filepath, image)
	logger.info(f"Overlay Masks: Saving {output_filepath}")
	else:
	label = cv2.convertScaleAbs(label, alpha=255.0 / label.max())
	contours, _ = cv2.findContours(label, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
	polygons.extend(self.to_polygons(contours))

	meta_json = {"image_size": image_size, "contours": len(polygons)}
	with open(os.path.join(output_dir, "meta.json"), "w") as fp:
	json.dump(meta_json, fp, indent=2)

	if output_contours:
	logger.info(f"Total Polygons: {len(polygons)}")
	with open(os.path.join(output_dir, "contours.json"), "w") as fp:
	json.dump({"count": len(polygons), "contours": polygons}, fp, indent=2)

	return d


	# Loss (adopted from Cellpose)
	class CellLoss:
	def __call__(self, y_pred, y):
	loss = 0.5 * F.mse_loss(y_pred[:, 1:], 5 * y[:, 1:]) + F.binary_cross_entropy_with_logits(
	y_pred[:, [0]], y[:, [0]]
	)
	return loss


	# Accuracy (adopted from Cellpose)
	class CellAcc:
	def __call__(self, mask_pred, mask_true):
	if isinstance(mask_true, torch.Tensor):
	mask_true = mask_true.cpu().numpy()

	if isinstance(mask_pred, torch.Tensor):
	mask_pred = mask_pred.cpu().numpy()

	# print("CellAcc mask_true", mask_true.shape, 'max', np.max(mask_true), ",
	# "'mask_pred', mask_pred.shape, 'max', np.max(mask_pred) )

	iou = _intersection_over_union(mask_true, mask_pred)[1:, 1:]
	tp = _true_positive(iou, th=0.5)

	fp = np.max(mask_pred) - tp
	fn = np.max(mask_true) - tp
	ap = tp / (tp + fp + fn)

	# print("CellAcc ap", ap, 'tp', tp, 'fp', fp, 'fn', fn)
	return ap