<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-
# Simple program for ITK image read/write in Python
# import itk


import numpy as np

from loguru import logger

import os.path

import re

import pydicom as dicom
dicom.config.debug(False)

import os.path as op
from typing import Union
from pathlib import Path
from . import rawN
from . import misc

# from sys import argv
from . import dcmtools
from . import image


def write(
    data3d: Union[np.ndarray, dict, image.DataPlus],
    path: Union[Path, str],
    filetype="auto",
    metadata: dict = None,
    allow_rescale_intercept: bool = False,
):
    """

    :param data3d: input ndarray or DataPlus
    :param path: output path, if braces are in the name ("dir/file{:04d}.dcm"), image stack is produced .
    Check function filename_format() for more details.
    :param filetype: dcm, png, h5, ... "image_stack"
    :param metadata: metadata f.e. {'voxelsize_mm': [3,2,2]}
    :param allow_rescale_intercept: SimpleITK does not allow to write 3D files with negative values int dicom. With this
    rescale_intercept option turned on, the unsigned data type will be used and rescale parameter in dicom metadata
    will be set.
    :return:
    """
    # if (type(data3d) == dict) or (type(data3d == image.DataPlus)):
    #     if metadata == None:
    #         metadata = data3d
    #         data3d = data3d.pop('data3d')
    #     else:
    #         logger.error("Cannot use first argument of type dict or DataPlus together with metadata.")
    #         raise ValueError("Cannot use first argument of type dict or DataPlus together with")

    dw = DataWriter()
    dw.Write3DData(
        data3d,
        path,
        filetype,
        metadata,
        allow_rescale_intercept=allow_rescale_intercept,
    )


class DataWriter:
    def __init__(self):
        self.stop_writing = False
        self.progress_callback = None

    def __get_segmentation_path(self, path):
        """ Create path with "_segmentation" suffix and keep extension.

        :param path:
        :return:
        """
        startpath, ext = os.path.splitext(path)
        segmentation_path = startpath + "_segmentation" + ext
        return segmentation_path

    def Write3DData(
        self,
        data3d,
        path,
        filetype="auto",
        metadata=None,
        progress_callback=None,
        sfin=True,
        allow_rescale_intercept=False,
    ):
        """
        :param data3d: input ndarray data
        :param path: output path, to specify slice number advanced formatting options (like {:06d}) can be used
        Check function filename_format() for more details.
        :param metadata: {'voxelsize_mm': [1, 1, 1]}
        :param filetype: dcm, vtk, rawiv, image_stack
        :param progress_callback: fuction for progressbar f.e. callback(value, minimum, maximum)
        :param sfin: Use separate file for segmentation if necessary
        :param allow_rescale_intercept: SimpleITK does not allow to write 3D files with negative values in dicom. With this
        rescale_intercept option turned on, the unsigned data type will be used and rescale parameter in dicom metadata
        will be set.


        """

        self.orig_path = path
        path = os.path.expanduser(path)

        tp = type(data3d)
        if tp in (image.DataPlus, np.ndarray):
            # copy
            datap = dict(data3d)
            data3d = datap.pop("data3d")
            metadata = datap

        if progress_callback is not None:
            self.progress_callback = progress_callback

        if filetype == "auto":
            startpath, ext = os.path.splitext(path)
            filetype = ext[1:].lower()

        segmentation = None
        if metadata is not None and "segmentation" in metadata.keys():
            segmentation_path = self.__get_segmentation_path(path)
            segmentation = metadata["segmentation"]

        mtd = {"voxelsize_mm": [1, 1, 1]}
        if metadata is not None:
            mtd.update(metadata)
        metadata = mtd

        if path.find("{") >= 0:
            filetype = "image_stack"
            # one_file_per_slice = True

        #     if one_file_per_slice:
        #         self._one_file_per_slice(self, data3d, path, filetype, metadata)
        #     else:
        #         self._all_in_one_file(self, data3d, path, filetype, metadata)
        #
        # def _all_in_one_file(self, data3d, path, filetype, metadata):

        if filetype in ["vtk", "tiff", "tif", "mhd", "nii", "raw"]:
            self._write_with_sitk(path, data3d, metadata, allow_rescale_intercept)<commit_msg>fixed wrong input type check<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-
# Simple program for ITK image read/write in Python
# import itk


import numpy as np

from loguru import logger

import os.path

import re

import pydicom as dicom
dicom.config.debug(False)

import os.path as op
from typing import Union
from pathlib import Path
from . import rawN
from . import misc

# from sys import argv
from . import dcmtools
from . import image


def write(
    data3d: Union[np.ndarray, dict, image.DataPlus],
    path: Union[Path, str],
    filetype="auto",
    metadata: dict = None,
    allow_rescale_intercept: bool = False,
):
    """

    :param data3d: input ndarray or DataPlus
    :param path: output path, if braces are in the name ("dir/file{:04d}.dcm"), image stack is produced .
    Check function filename_format() for more details.
    :param filetype: dcm, png, h5, ... "image_stack"
    :param metadata: metadata f.e. {'voxelsize_mm': [3,2,2]}
    :param allow_rescale_intercept: SimpleITK does not allow to write 3D files with negative values int dicom. With this
    rescale_intercept option turned on, the unsigned data type will be used and rescale parameter in dicom metadata
    will be set.
    :return:
    """

    dw = DataWriter()
    dw.Write3DData(
        data3d,
        path,
        filetype,
        metadata,
        allow_rescale_intercept=allow_rescale_intercept,
    )


class DataWriter:
    def __init__(self):
        self.stop_writing = False
        self.progress_callback = None

    def __get_segmentation_path(self, path):
        """ Create path with "_segmentation" suffix and keep extension.

        :param path:
        :return:
        """
        startpath, ext = os.path.splitext(path)
        segmentation_path = startpath + "_segmentation" + ext
        return segmentation_path

    def Write3DData(
        self,
        data3d,
        path,
        filetype="auto",
        metadata=None,
        progress_callback=None,
        sfin=True,
        allow_rescale_intercept=False,
    ):
        """
        :param data3d: input ndarray data
        :param path: output path, to specify slice number advanced formatting options (like {:06d}) can be used
        Check function filename_format() for more details.
        :param metadata: {'voxelsize_mm': [1, 1, 1]}
        :param filetype: dcm, vtk, rawiv, image_stack
        :param progress_callback: fuction for progressbar f.e. callback(value, minimum, maximum)
        :param sfin: Use separate file for segmentation if necessary
        :param allow_rescale_intercept: SimpleITK does not allow to write 3D files with negative values in dicom. With this
        rescale_intercept option turned on, the unsigned data type will be used and rescale parameter in dicom metadata
        will be set.


        """

        self.orig_path = path
        path = os.path.expanduser(path)

        tp = type(data3d)
        if tp in (image.DataPlus, dict):
            # copy
            datap = dict(data3d)
            data3d = datap.pop("data3d")
            metadata = datap
        elif tp == np.ndarray:
            pass
        else:
            logger.warning(f"Type '{tp}' may not be compatible with datawriter.")

        if progress_callback is not None:
            self.progress_callback = progress_callback

        if filetype == "auto":
            startpath, ext = os.path.splitext(path)
            filetype = ext[1:].lower()

        segmentation = None
        if metadata is not None and "segmentation" in metadata.keys():
            segmentation_path = self.__get_segmentation_path(path)
            segmentation = metadata["segmentation"]

        mtd = {"voxelsize_mm": [1, 1, 1]}
        if metadata is not None:
            mtd.update(metadata)
        metadata = mtd

        if path.find("{") >= 0:
            filetype = "image_stack"
            # one_file_per_slice = True

        #     if one_file_per_slice:
        #         self._one_file_per_slice(self, data3d, path, filetype, metadata)
        #     else:
        #         self._all_in_one_file(self, data3d, path, filetype, metadata)
        #
        # def _all_in_one_file(self, data3d, path, filetype, metadata):

        if filetype in ["vtk", "tiff", "tif", "mhd", "nii", "raw"]:
            self._write_with_sitk(path, data3d, metadata, allow_rescale_intercept)
{'fraction_non_alphanumeric': 0.08978775419491054, 'fraction_numerical': 0.008223135903989333, 'mean_word_length': 2.3975084937712343, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4395671', 'n_tokens_mistral': 2716, 'n_tokens_neox': 2648, 'n_words': 846}