more tests

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
 

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
         )

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(

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):

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])