diff --git a/examples/infer_contrasting_ct_dicom_brain.sh b/examples/infer_contrasting_ct_dicom_brain.sh index c2145b8a..dc06de8e 100755 --- a/examples/infer_contrasting_ct_dicom_brain.sh +++ b/examples/infer_contrasting_ct_dicom_brain.sh @@ -2,7 +2,7 @@ data_path=$1 out_path=$2 if [ -z "$data_path" ]; then - data_path="../innofw/data/rtk/infer/" + data_path="https://api.blackhole.ai.innopolis.university/public-datasets/rtk/infer.zip" echo "Using default data path $data_path" fi diff --git a/innofw/core/datamodules/lightning_datamodules/coco_rtk.py b/innofw/core/datamodules/lightning_datamodules/coco_rtk.py index 3337ba50..a5e240b9 100644 --- a/innofw/core/datamodules/lightning_datamodules/coco_rtk.py +++ b/innofw/core/datamodules/lightning_datamodules/coco_rtk.py @@ -18,6 +18,15 @@ def __call__(self, image, **kwargs): return image +DEFAULT_TRANSFORM = albu.Compose( + [ + albu.Resize(256, 256), + albu.Lambda(image=CustomNormalize()), + ToTensorV2(transpose_mask=True), + ] +) + + class DicomCocoComplexingDataModule(BaseLightningDataModule): task = ["image-detection", "image-segmentation"] dataset = DicomCocoDatasetRTK @@ -59,14 +68,7 @@ def setup_infer(self): if self.aug: transform = Augmentation(self.aug["test"]) else: - - transform = albu.Compose( - [ - albu.Resize(256, 256), - albu.Lambda(image=CustomNormalize()), - ToTensorV2(transpose_mask=True), - ] - ) + transform = DEFAULT_TRANSFORM if str(self.predict_source).split("/")[-1] in ["mrt", "ct"]: self.predict_source = self.predict_source.parent cont = os.listdir(self.predict_source) @@ -152,14 +154,7 @@ def setup_infer(self): if self.aug: transform = Augmentation(self.aug["test"]) else: - - transform = albu.Compose( - [ - albu.Resize(256, 256), - albu.Lambda(image=CustomNormalize()), - ToTensorV2(transpose_mask=True), - ] - ) + transform = DEFAULT_TRANSFORM self.predict_dataset = self.dataset( data_dir=str(self.predict_source), transform=transform ) diff --git a/innofw/core/datasets/coco_rtk.py b/innofw/core/datasets/coco_rtk.py index d06b9b07..7d092021 100644 --- a/innofw/core/datasets/coco_rtk.py +++ b/innofw/core/datasets/coco_rtk.py @@ -83,7 +83,6 @@ def extract_digits(s): if dicom_path.endswith(img["file_name"]): new_images += [img] self.images = new_images - self.images.sort(key=lambda x: extract_digits(x["file_name"])) else: self.dicom_paths.sort() @@ -95,6 +94,11 @@ def __len__(self): return len(self.dicom_paths) def get_dicom(self, i): + """no annotation data + + :param i: + :return: + """ dicom_path = self.dicom_paths[i] dicom = pydicom.dcmread(dicom_path) @@ -171,9 +175,7 @@ def get_mask(self, anns, image_info): ) for ann in anns: segmentation = ann["segmentation"] - category_id = ( - ann["category_id"] - 1 - ) # Приведение category_id к индексу слоя + category_id = ann["category_id"] - 1 if isinstance(segmentation, list): # полигональная аннотация for polygon in segmentation: poly_mask = self._polygon_to_mask( diff --git a/innofw/utils/data_utils/preprocessing/CT_hemorrhage_contrast_rtk.py b/innofw/utils/data_utils/preprocessing/CT_hemorrhage_contrast_rtk.py index 42e5313b..5d63c07a 100644 --- a/innofw/utils/data_utils/preprocessing/CT_hemorrhage_contrast_rtk.py +++ b/innofw/utils/data_utils/preprocessing/CT_hemorrhage_contrast_rtk.py @@ -1,12 +1,16 @@ from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser from pathlib import Path +from urllib.parse import urlparse import os from tqdm import tqdm import cv2 import numpy as np -from innofw.core.datasets.coco_rtk import DicomCocoDatasetRTK + +from innofw.core.datamodules.lightning_datamodules.coco_rtk import ( + DicomCocoDataModuleRTK, +) from innofw.utils.data_utils.rtk.CT_hemorrhage_metrics import transform @@ -18,14 +22,22 @@ def hemorrhage_contrast(input_path: str, output_folder: str = None): output_folder = str(output_folder) except TypeError: raise ValueError(f"Wrong path to save: {output_folder}") + if urlparse(input_path): + default_path = "innofw/data/rtk/infer/" + path = {"source": input_path, "target": default_path} + else: + path = {"source": input_path, "target": input_path} + dm = DicomCocoDataModuleRTK(infer=path, transform=transform) + dm.setup_infer() + dataloader = dm.predict_dataloader() + dataset = dataloader.dataset + dataset.transform = transform - dataset = DicomCocoDatasetRTK(data_dir=input_path, transform=transform) if len(dataset) == 0: raise Warning(f"empty dataset with the directory {input_path}") else: for x in (pbar := tqdm(dataset)): path = x["path"] - mask = x.get("mask", None) contrasted_image = x["image"] raw_image = x.get("raw_image", None) @@ -34,32 +46,13 @@ def hemorrhage_contrast(input_path: str, output_folder: str = None): os.makedirs(output_folder, exist_ok=True) output_path = os.path.join(output_folder, basename + "_raw.npy") - try: - np.save(output_path, raw_image) - except: - pbar.set_description(f"wrong path {output_path}") + np.save(output_path, raw_image) output_path = os.path.join(output_folder, basename + "_mask.png") - try: - if cv2.imwrite(output_path, mask): - pbar.set_description(f"could not write mask as {output_path}") - else: - raise Exception - except: - pbar.set_description(f"could not write mask as {output_path}") + cv2.imwrite(output_path, mask) output_path = os.path.join(output_folder, basename + "_image.png") - try: - if cv2.imwrite(output_path, contrasted_image): - pbar.set_description( - f"could not write contrasted image as {output_path}" - ) - else: - raise Exception - except: - pbar.set_description( - f"could not write contrasted image as {output_path}" - ) + cv2.imwrite(output_path, contrasted_image) def callback(arguments): diff --git a/tests/unit/datamodules/lightning_datamodules/test_rtk.py b/tests/unit/datamodules/lightning_datamodules/test_rtk.py index 034bfdc4..f78bb9eb 100644 --- a/tests/unit/datamodules/lightning_datamodules/test_rtk.py +++ b/tests/unit/datamodules/lightning_datamodules/test_rtk.py @@ -2,63 +2,87 @@ import shutil import pytest +import numpy as np from innofw.core.datamodules.lightning_datamodules.coco_rtk import ( + DEFAULT_TRANSFORM, DicomCocoComplexingDataModule, DicomCocoDataModuleRTK, ) -from innofw.core.datasets.coco_rtk import DicomCocoDatasetRTK from innofw.core.datamodules.pandas_datamodules.lung_description_decision_datamodule import ( LungDescriptionDecisionPandasDataModule, ) +from innofw.core.datasets.coco_rtk import DicomCocoDatasetRTK +from innofw.utils.data_utils.preprocessing.CT_hemorrhage_contrast_rtk import ( + hemorrhage_contrast, +) rtk_complex = "https://api.blackhole.ai.innopolis.university/public-datasets/rtk/complex_infer.zip" rtk_segm = "https://api.blackhole.ai.innopolis.university/public-datasets/rtk/infer.zip" lungs = "https://api.blackhole.ai.innopolis.university/public-datasets/rtk/labels.zip" -def test_DicomCocoComplexingDataModule(): +transforms = [None, DEFAULT_TRANSFORM] + + +@pytest.mark.parametrize("transform", transforms) +def test_DicomCocoComplexingDataModule(transform): target_dir = "./data/complex/infer" if os.path.exists(target_dir): shutil.rmtree(target_dir) path = {"source": rtk_complex, "target": target_dir} - dm = DicomCocoComplexingDataModule(infer=path) + dm = DicomCocoComplexingDataModule(infer=path, transform=transform) dm.setup_infer() ds = dm.predict_dataloader() for batch in ds: break for k in ["image", "mask", "path"]: - assert k in batch + assert k in batch, f"no suck key {k}" -def test_DicomCocoDataModuleRTK(): +@pytest.mark.parametrize("transform", transforms) +def test_DicomCocoDataModuleRTK(transform): target_dir = "./data/rtk/infer" if os.path.exists(target_dir): shutil.rmtree(target_dir) path = {"source": rtk_segm, "target": target_dir} - dm = DicomCocoDataModuleRTK(infer=path) + dm = DicomCocoDataModuleRTK(infer=path, transform=transform) dm.setup_infer() ds = dm.predict_dataloader() for batch in ds: break for k in ["image", "path"]: - assert k in batch + assert k in batch, f"no suck key {k}" -def test_DicomCocoDataset_rtk(): +@pytest.mark.parametrize("transform", transforms) +def test_DicomCocoDataset_rtk(transform): """ - import this test to run after previous + import this test to run after test_DicomCocoDataModuleRTK """ path = "./data/rtk/infer" - ds = DicomCocoDatasetRTK(data_dir=path) + ds = DicomCocoDatasetRTK(data_dir=path, transform=transform) for batch in ds: break for k in ["image", "mask", "path"]: - assert k in batch + assert k in batch, f"no suck key {k}" -# +@pytest.mark.parametrize("transform", transforms) +def test_DicomCocoDataset_rtk_no_annotation(transform): + """ + import this test to run after previous + """ + path = "./data/rtk/infer" + annotations_file = os.path.join(path, "annotations", "instances_default.json") + if os.path.exists(annotations_file): + os.remove(annotations_file) + ds = DicomCocoDatasetRTK(data_dir=path, transform=transform) + for batch in ds: + break + for k in ["image", "path"]: + assert k in batch, f"no suck key {k}" def test_datamodule_description(): @@ -71,3 +95,16 @@ def test_datamodule_description(): ds = dm.predict_dataloader() for key in "x", "y": assert key in ds + + +def test_hemor_contrast(tmp_path_factory): + target_dir = "./data/rtk/infer" + if os.path.exists(target_dir): + shutil.rmtree(target_dir) + out_ = tmp_path_factory.mktemp("out") + hemorrhage_contrast(input_path=rtk_segm, output_folder=out_) + content = os.listdir(out_) + assert len(content) > 0 + assert len(content) % 3 == 0 + assert np.any([x.endswith("npy") for x in content]) + assert np.any([x.endswith("png") for x in content])