dicomqr.py

"""
DICOM-QR utilities.

Depends on pynetdicom3: https://github.com/scaramallion/pynetdicom3
"""
import json
from typing import Dict, Optional

import click
import structlog
from pydicom.dataset import Dataset
from pynetdicom3 import AE, QueryRetrieveSOPClassList
from pynetdicom3.sop_class import QueryRetrieveFindServiceClass, VerificationSOPClass
from pynetdicom3.status import code_to_category
from structlog import BoundLogger


class DicomException(BaseException):
    """Base class for dicom exceptions."""
    def __init__(self, message):
        self.message = message


class DicomAssociate:
    """Context manager to associate with a dicom peer.

    Implements DICOM association as a context manager that
    returns an instance to a pynetdicom3.AE.

    Args:
        host: host of the dicom peer to associate with
        port: port of the dicom peer to associate with
        remote_aet: AE title of the dicom peer to associate with
        local_aet: local AE title
        do_echo_check: if True, do an echo call to the remote host

    Returns:
        assoc: Association instance
    """
    def __init__(self, host: str, port: int, remote_aet: str, local_aet: str, do_echo_check: bool=False) \
            -> None:
        self.host = host
        self.port = port
        self.remote_aet = remote_aet
        self.local_aet = local_aet
        self.do_echo_check = do_echo_check
        self.assoc = None

    def __enter__(self):
        scu_sop_classes = QueryRetrieveSOPClassList + [VerificationSOPClass]
        ae = AE(scu_sop_class=scu_sop_classes, ae_title=self.local_aet)
        self.assoc = ae.associate(self.host, self.port, ae_title=self.remote_aet)
        if not self.assoc.is_established:
            raise DicomException(
                f"Failed to associate {self.remote_aet}@{self.host}:{self.port}: {self.assoc}")
        elif self.do_echo_check:
            try:
                status = self.assoc.send_c_echo()
            except ValueError as err:
                raise DicomException(
                    f"Echo failure {self.remote_aet}@{self.host}:{self.port}: {self.assoc}. Error: {err}")
            if status:
                success_status_category = 'Success'
                status_category = code_to_category(status.Status)
                if not status_category == success_status_category:
                    raise DicomException(
                        f"Echo failure {self.remote_aet}@{self.host}:{self.port}: {self.assoc}. Status: "
                        f"{status_category}")
        return self.assoc

    def __exit__(self, *args):
        if self.assoc:
            self.assoc.release()


class DicomQR:
    """
    DicomQR class to make Query/Retrieve calls.

    Args:
        host: the host/ip of the remote DICOM server
        port: the port on which the remote DICOM server listens
        remote_aet: the application entity title of the remote server to use when querying
        local_aet: local application entity title
        log: BoundLogger instance (optional)
    """
    def __init__(
        self, host: str, port: int, remote_aet: str, local_aet: str, log: Optional[BoundLogger] = None) \
            -> None:
        self.host = host
        self.port = port
        self.remote_aet = remote_aet
        self.local_aet = local_aet
        self.log = log

    def _handle_qr_find_responses(self, responses, handler):
        """Helper to iterate over DICOM responses.

        When the response has data the handler is called.

        Possible status categories" 'Cancel', 'Failure', 'Pending', 'Success', 'Warning'
        """
        for (status, ds) in responses:
            if QueryRetrieveFindServiceClass.statuses[status.Status][0] == 'Success':
                return
            elif QueryRetrieveFindServiceClass.statuses[status.Status][0] == 'Pending':
                handler(ds)
            else:
                raise DicomException(f"Failed to c_find: {status}")

    def _query_series_images(self, select: Dict) -> Dict:
        """Actual retrieval of the number of series images based on the select query items.

        Args:
            select: dictionary with query items. Example select dictionaries:

        Returns:
            number of images each series contain for the given select parameters
        """
        series = {}
        with DicomAssociate(self.host, self.port, self.remote_aet, self.local_aet) as assoc:
            # retrieve all series in the study
            ds = Dataset()
            ds.QueryRetrieveLevel = 'SERIES'
            for tag, value in select.items():
                setattr(ds, tag, value)
            ds.SeriesInstanceUID = ''
            responses = assoc.send_c_find(ds, query_model='S')

            def handler(ds):
                series[ds.SeriesInstanceUID] = 0

            self._handle_qr_find_responses(responses, handler)

            # for all series, count the number of images
            for s in series:
                ds = Dataset()
                ds.QueryRetrieveLevel = 'IMAGE'
                for tag, value in select.items():
                    setattr(ds, tag, value)
                ds.SeriesInstanceUID = s
                ds.SOPInstanceUID = ''
                responses = assoc.send_c_find(ds, query_model='S')

                def handler(ds):
                    series[ds.SeriesInstanceUID] = series[ds.SeriesInstanceUID] + 1

                self._handle_qr_find_responses(responses, handler)
        return series

    def series_images(self, select: Dict) -> Dict:
        """
        Return a dictionary containing the number of images of each series belonging to the studies that are
        returned by the select query items.

        Args:
            select: dictionary with query items. Example select dictionaries:
                - {'AccessionNumber': '1574164977304'}
                - {'AccessionNumber': '9196154216199', 'PatientID': '123456'}
                - {'StudyInstanceUID': '1.3.6.1.4.1.14519.5.2.1.7009.9004.160416029404301360186441590'}

        Returns:
            number of images each series contain for the given select parameters
        """
        try:
            return self._query_series_images(select)
        except DicomException:
            if self.log:
                self.log.exception(f"Failed to get number of series images with QR.")
            return {}

    def series_images_by_study_uid(self, study_uid: str) -> Dict:
        """Helper function that returns a dictionary containing the number of images of each series belonging to the
        study with the given study_uid.

        Args:
            study_uid: study instance uid

        Returns:
            number of images each series contain for the given study
        """
        return self.series_images({'StudyInstanceUID': study_uid})


@click.command()
@click.option('-h', '--host', default='localhost', type=str, help='Remote host (default: localhost)')
@click.option('-p', '--port', default=4242, type=int, help="Remote port (default: 4242)")
@click.option('--remote-aet', default='ORTHANC', type=str, help="Remote AE title (default: ORTHANC)")
@click.option('--local-aet', default='PYNETDICOM', type=str, help="Local AE title (default: PYNETDICOM)")
@click.argument('study_uid', type=str)
@click.option('-v', '--verbose', is_flag=True, help='Output more logging.')
def series_images_cli(host: str, port: int, remote_aet: str, local_aet: str, study_uid: str, verbose: bool):
    """
    CLI for testing series_images_by_study_uid query.

    Usage:

    Orthanc
    -------
    Make sure PYNETDICOM is added to /etc/orthanc/orthanc.json. Example configuration:
    "DicomModalities" : {
      "findscu" : [ "FINDSCU", "127.0.0.1", 1234 ],
      "pynetdicom": ["PYNETDICOM", "127.0.0.1", 1234],
    }
    $ python dicomqr.py <study_uid>

    """
    log = structlog.get_logger(__name__) if verbose else None
    dicom_qr = DicomQR(host, port, remote_aet, local_aet, log)
    series_images = dicom_qr.series_images_by_study_uid(study_uid)
    if verbose:
        print(series_images)
    else:
        # show as json to make it parseable by other programs
        print(json.dumps(series_images))


if __name__ == "__main__":
    import logging
    logger = logging.getLogger('pynetdicom3')
    logger.setLevel(logging.DEBUG)
    series_images_cli()