dicom-transfer.py

#!/usr/bin/env python3
"""
DICOM Transfer Script

This script transfers medical imaging files (.dcm, .nifi, .nrrd, .IMA) from a local directory
to a destination DICOM server using the DICOM C-STORE protocol.

Configuration Parameters:
- DESTINATION_SERVER: IP address or hostname of the destination DICOM server
- DESTINATION_SERVER_PORT: Port number of the destination DICOM server
- DESTINATION_SERVER_AET: Application Entity Title of the destination server
- LOCAL_AET: Local Application Entity Title
- INPUT_DIRECTORY: Path to the directory containing files to transfer
"""

import os
import sys
import argparse
import logging
from pathlib import Path
from typing import List, Dict, Any
import time
import threading
import queue


try:
    from pydicom import dcmread, Dataset
    from pydicom.errors import InvalidDicomError
    from pynetdicom import AE, StoragePresentationContexts
    # from pynetdicom.sop_class import CTImageStorage, MRImageStorage, PETImageStorage
    # pynetdicom.sop_class.LegacyConvertedEnhancedPETImageStorage
    from pynetdicom.sop_class import (
        CTImageStorage,
        MRImageStorage,
        PositronEmissionTomographyImageStorage,
        RTImageStorage,
        RTDoseStorage,
        RTStructureSetStorage,
        RTPlanStorage,
        RTBeamsTreatmentRecordStorage,
        RTBrachyTreatmentRecordStorage,
        RTIonBeamsTreatmentRecordStorage,
        EnhancedCTImageStorage,
        EnhancedPETImageStorage,
        SecondaryCaptureImageStorage,
        ComputedRadiographyImageStorage,
        DigitalXRayImageStorageForPresentation,
        DigitalXRayImageStorageForProcessing,
        UltrasoundImageStorage,
        EnhancedUSVolumeStorage,
    )
    from pynetdicom.events import Event
    from pynetdicom.status import Status
except ImportError as e:
    print(f"Error: Required DICOM libraries not found. Please install with: pip install pydicom pynetdicom")
    print(f"Missing: {e}")
    sys.exit(1)


class DicomTransfer:
    """DICOM file transfer class for sending medical imaging files to a DICOM server."""
    def __init__(self, config: Dict[str, Any]):
        """
        Initialize the DICOM transfer with configuration parameters.
        
        Args:
            config: Dictionary containing transfer configuration
        """
        self.config = config
        self.setup_logging()
        
        # Thread-safe queue for found files
        self.found_files = queue.Queue()
        
        # Thread-safe statistics
        self.stats_lock = threading.Lock()
        self.stats = {
            'total_files': 0,
            'successful_transfers': 0,
            'failed_transfers': 0,
            'skipped_files': 0
        }
        
        # Control flags
        self.finding_complete = threading.Event()
        self.shutdown = threading.Event()
        self.worker_threads = []
        
        # Number of transfer worker threads
        self.num_workers = self.config.get('num_workers', 5)
    
    def setup_logging(self):
        """Setup logging configuration."""
        log_level = self.config.get('log_level', 'INFO').upper()
        logging.basicConfig(
            level=getattr(logging, log_level),
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler('dicom_transfer.log'),
                logging.StreamHandler(sys.stdout)
            ]
        )
        self.logger = logging.getLogger(__name__)
    
    def find_files(self, directory: str, is_root_call: bool = True):
        """
        Recursively find all supported files in the directory and add them to the queue.
        This method runs in a separate thread and preserves the original recursive pattern.
        
        Args:
            directory: Path to search for files
            is_root_call: True if this is the initial (root) call, False for recursive calls
        """
        input_path = Path(directory)
        if not input_path.exists():
            self.logger.error(f"Input directory does not exist: {directory}")
            if is_root_call:
                self.finding_complete.set()
            return
        
        if not input_path.is_dir():
            self.logger.error(f"Input path is not a directory: {directory}")
            if is_root_call:
                self.finding_complete.set()
            return
        
        self.logger.debug(f"Scanning directory: {directory}")
        
        try:
            # Use os.walk() but process only one level, then manually recurse (as in original)
            # This preserves the original pattern even though os.walk() already recurses
            for root, dirs, files in os.walk(input_path):
                # Process files in current directory
                for file in files:
                    file_path = Path(root) / file
                    if file_path.suffix in self.config.get('supported_extensions', ['.dcm', '.nifi', '.nrrd', '.IMA']):
                        self.found_files.put(file_path)
                        with self.stats_lock:
                            self.stats['total_files'] += 1
                
                # Recursively call find_files for each subdirectory (as in original pattern)
                for new_dir in dirs:
                    self.find_files(Path(root) / new_dir, is_root_call=False)
                
                # Break after first iteration since we handle recursion manually
                # (This prevents os.walk() from automatically recursing)
                break
        except Exception as e:
            self.logger.error(f"Error during file discovery in {directory}: {e}")
        
        # Only signal completion if this was the root call
        if is_root_call:
            self.finding_complete.set()
            self.logger.info(f"File discovery completed. Found {self.stats['total_files']} files total")
    
    def setup_dicom_association(self):
        """Setup DICOM association with the destination server. Returns (ae, assoc) tuple."""
        try:
            # Create Application Entity
            ae = AE(ae_title=self.config['local_aet'])
            common_sop_classes = [
                CTImageStorage,
                MRImageStorage,
                PositronEmissionTomographyImageStorage,
                EnhancedCTImageStorage,
                EnhancedPETImageStorage,
                SecondaryCaptureImageStorage,
                ComputedRadiographyImageStorage,
                DigitalXRayImageStorageForPresentation,
                DigitalXRayImageStorageForProcessing,
                UltrasoundImageStorage,
                RTStructureSetStorage,  # Added RTSTRUCT
            ]
            transfer_syntaxes = [
                '1.2.840.10008.1.2.1',   # Explicit VR Little Endian (preferred)
                '1.2.840.10008.1.2',     # Implicit VR Little Endian
                '1.2.840.10008.1.2.2'    # Explicit VR Big Endian
            ]
            for sop_class in common_sop_classes:
                ae.add_requested_context(sop_class, transfer_syntax=transfer_syntaxes)
            # Use StoragePresentationContexts which includes all standard transfer syntaxes
            # This is the recommended way for C-STORE operations
            # ae.requested_contexts = StoragePresentationContexts
            
            assoc = ae.associate(
                addr=self.config['destination_server'],
                port=self.config['destination_port'],
                ae_title=self.config['destination_aet']
            )
            
            if assoc and assoc.is_established:
                accepted = [ pc for pc in assoc.accepted_contexts]
                self.logger.debug(f"DICOM association established successfully with {len(accepted)} accepted presentation contexts")
                return (ae, assoc)
            else: 
                self.logger.error("Failed to establish DICOM association")
                if assoc:
                    # Log rejection details
                    if hasattr(assoc, 'accepted_contexts'):
                        self.logger.debug(f"Accepted contexts: {len(assoc.accepted_contexts)}")
                    if hasattr(assoc, 'rejected_contexts'):
                        rejected = assoc.rejected_contexts
                        if rejected:
                            self.logger.error(f"Rejected contexts: {len(rejected)}")
                            for rc in rejected:
                                if hasattr(rc, 'abstract_syntax'):
                                    self.logger.error(f"Rejected SOP Class: {rc.abstract_syntax}")
                                if hasattr(rc, 'result'):
                                    self.logger.error(f"  Result: {rc.result}")
                                if hasattr(rc, 'reason'):
                                    self.logger.error(f"  Reason: {rc.reason}")
                    # Check if association was rejected entirely
                    if hasattr(assoc, 'release'):
                        # Try to get more details
                        self.logger.error("Association rejection - check AET configuration on dcm4chee server")
                        self.logger.error(f"Local AET '{self.config['local_aet']}' may not be authorized to connect to '{self.config['destination_aet']}'")
                return (None, None)

        except Exception as e:
            self.logger.error(f"Error setting up DICOM association: {e}")
            return (None, None)
            

    def close_association(self, assoc):
        """Close the DICOM association."""
        if assoc and assoc.is_established:
            assoc.release()
            self.logger.debug("DICOM association released")
    
    def is_dicom_file(self, file_path: Path) -> bool:
        """
        Check if a file is a valid DICOM file.
        
        Args:
            file_path: Path to the file to check
            
        Returns:
            True if file is valid DICOM, False otherwise
        """
        try:
            supported_ext = self.config.get('supported_extensions', ['.dcm', '.nifi', '.nrrd', '.IMA'])
            if file_path.suffix in supported_ext:
                # Try to read as DICOM
                dcmread(file_path, stop_before_pixels=True)
                return True
            else:
                # For .nifi and .nrrd files, we'll assume they need conversion
                # or special handling - for now, we'll skip them with a warning
                self.logger.warning(f"Not supported file format detected: {file_path.suffix}")
                return False
        except InvalidDicomError:
            self.logger.warning(f"Invalid DICOM file: {file_path}")
            return False
        except Exception as e:
            self.logger.error(f"Error reading file {file_path}: {e}")
            return False
    
    def transfer_file(self, file_path: Path, assoc) -> bool:
        """
        Transfer a single file to the DICOM server.
        
        Args:
            file_path: Path to the file to transfer
            assoc: DICOM association object
            
        Returns:
            True if transfer successful, False otherwise
        """
        try:
            if not self.is_dicom_file(file_path):
                with self.stats_lock:
                    self.stats['skipped_files'] += 1
                return False
            
            # Read DICOM file
            ds = dcmread(file_path)
            
            # Send C-STORE request
            status = assoc.send_c_store(ds)
            
            # Get numeric status code
            # Status can be a status object or numeric value
            if hasattr(status, 'Status'):
                status_code = status.Status
            elif hasattr(status, 'status'):
                status_code = status.status
            else:
                # Status is already a numeric value
                status_code = status
            
            # Convert to integer if needed
            if isinstance(status_code, (tuple, list)) and len(status_code) > 0:
                status_code = status_code[0]
            
            # DICOM C-STORE status codes:
            # 0x0000-0xBFFF: Success (with or without warnings)
            # 0xC000-0xFFFF: Failure
            # We consider anything < 0xC000 as success (including replacements/warnings)
            if isinstance(status_code, int) and status_code < 0xC000:
                # Success (including warnings like file replacement)
                status_msg = f"Status {hex(status_code)}"
                if status_code == 0x0000:
                    self.logger.info(f"Successfully transferred: {file_path.name}")
                elif status_code == 0xB000:
                    self.logger.info(f"Successfully transferred (with coercion): {file_path.name} - Status {hex(status_code)}")
                elif status_code == 0xB007:
                    self.logger.info(f"Successfully transferred (coerced to match SOP): {file_path.name} - Status {hex(status_code)}")
                else:
                    self.logger.info(f"Successfully transferred (with warning): {file_path.name} - Status {hex(status_code)}")
                
                with self.stats_lock:
                    self.stats['successful_transfers'] += 1
                return True
            else:
                # Failure
                self.logger.error(f"Transfer failed for {file_path.name}: Status {hex(status_code) if isinstance(status_code, int) else status_code}")
                with self.stats_lock:
                    self.stats['failed_transfers'] += 1
                return False
                
        except Exception as e:
            self.logger.error(f"Error transferring {file_path}: {e}")
            with self.stats_lock:
                self.stats['failed_transfers'] += 1
            return False
    
    def transfer_worker(self, worker_id: int):
        """
        Worker thread that processes files from the queue.
        
        Args:
            worker_id: Unique identifier for this worker thread
        """
        self.logger.info(f"Transfer worker {worker_id} started")
        
        # Setup DICOM association for this worker
        ae, assoc = self.setup_dicom_association()
        if not assoc or not assoc.is_established:
            self.logger.error(f"Worker {worker_id}: Failed to establish DICOM association. Exiting.")
            return
        
        try:
            while not self.shutdown.is_set():
                try:
                    # Try to get a file from the queue with timeout
                    # This allows us to periodically check shutdown flag
                    file_path = self.found_files.get(timeout=2.0)
                    
                    # Process the file
                    self.logger.debug(f"Worker {worker_id} processing: {file_path.name}")
                    self.transfer_file(file_path, assoc)
                    self.found_files.task_done()
                    
                except queue.Empty:
                    # Queue is empty, check if finding is complete
                    if self.finding_complete.is_set() and self.found_files.empty():
                        # Finding is complete and queue is empty, we're done
                        self.logger.debug(f"Worker {worker_id}: Queue empty and finding complete, exiting")
                        break
                    else:
                        # Still finding files, sleep for 2 seconds as requested
                        self.logger.debug(f"Worker {worker_id}: Queue empty, sleeping 2 seconds")
                        time.sleep(2)
                        
        except Exception as e:
            self.logger.error(f"Worker {worker_id} error: {e}")
        finally:
            # Close association if it was established
            if assoc:
                self.close_association(assoc)
            self.logger.info(f"Transfer worker {worker_id} finished")
    
    def start_transfer(self, directory: str):
        """
        Start the concurrent transfer process.
        
        Args:
            directory: Directory to scan for files
        """
        # Start file finding thread (is_root_call=True is the default)
        finder_thread = threading.Thread(target=self.find_files, args=(directory, True), daemon=False)
        finder_thread.start()
        
        # Start transfer worker threads
        self.logger.info(f"Starting {self.num_workers} transfer worker threads...")
        for i in range(self.num_workers):
            worker_thread = threading.Thread(target=self.transfer_worker, args=(i+1,), daemon=False)
            worker_thread.start()
            self.worker_threads.append(worker_thread)
        
        # Wait for finder to complete
        finder_thread.join()
        self.logger.info("File discovery completed")
        
        # Wait for all workers to complete
        self.logger.info("Waiting for all transfer workers to complete...")
        for worker_thread in self.worker_threads:
            worker_thread.join()
        
        self.logger.info("All transfers completed")
    
    def stop(self):
        """Signal all threads to stop."""
        self.shutdown.set()
        # Wait for all workers to finish
        for worker_thread in self.worker_threads:
            worker_thread.join(timeout=5)
    
    def transfer_files(self, files: List[Path]):
        """
        Transfer all files to the DICOM server.
        
        Args:
            files: List of file paths to transfer
        """
        if not self.setup_dicom_association():
            self.logger.error("Cannot establish DICOM association. Aborting transfer.")
            return
        
        try:
            self.logger.info(f"Starting transfer of {len(files)} files...")
            
            for i, file_path in enumerate(files, 1):
                self.logger.info(f"Processing file {i}/{len(files)}: {file_path.name}")
                
                success = self.transfer_file(file_path)
                
                # Add small delay between transfers to avoid overwhelming the server
                if i < len(files):
                    time.sleep(0.1)
            
            self.logger.info("Transfer process completed")
            
        finally:
            self.close_association()
    
    def print_statistics(self):
        """Print transfer statistics."""
        with self.stats_lock:
            stats = self.stats.copy()
        
        print("\n" + "="*50)
        print("TRANSFER STATISTICS")
        print("="*50)
        print(f"Total files found: {stats['total_files']}")
        print(f"Successfully transferred: {stats['successful_transfers']}")
        print(f"Failed transfers: {stats['failed_transfers']}")
        print(f"Skipped files: {stats['skipped_files']}")
        
        if stats['total_files'] > 0:
            success_rate = (stats['successful_transfers'] / stats['total_files']) * 100
            print(f"Success rate: {success_rate:.1f}%")
        print("="*50)


def parse_arguments():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Transfer medical imaging files to a DICOM server",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python dicom-transfer.py --server 192.168.1.100 --port 104 --remote-aet PACS --local-aet CLIENT --input /path/to/images
  python dicom-transfer.py --config config.json
        """
    )
    
    # Configuration options
    parser.add_argument('--server', '-s', 
                       help='Destination DICOM server IP address or hostname')
    parser.add_argument('--port', '-p', type=int,
                       help='Destination DICOM server port')
    parser.add_argument('--remote-aet', '-r',
                       help='Destination server Application Entity Title')
    parser.add_argument('--local-aet', '-l',
                       help='Local Application Entity Title')
    parser.add_argument('--input', '-i',
                       help='Input directory containing files to transfer')
    
    # Optional parameters
    parser.add_argument('--config', '-c',
                       help='Configuration file (JSON format)')
    parser.add_argument('--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'],
                       default='INFO', help='Logging level')
    parser.add_argument('--dry-run', action='store_true',
                       help='Show what would be transferred without actually transferring')
    
    return parser.parse_args()


def load_config_file(config_path: str) -> Dict[str, Any]:
    """
    Load configuration from JSON file.
    
    Args:
        config_path: Path to configuration file
        
    Returns:
        Configuration dictionary
    """
    import json
    
    try:
        with open(config_path, 'r') as f:
            config = json.load(f)
        return config
    except FileNotFoundError:
        print(f"Error: Configuration file not found: {config_path}")
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"Error: Invalid JSON in configuration file: {e}")
        sys.exit(1)


def validate_config(config: Dict[str, Any]) -> bool:
    """
    Validate configuration parameters.
    
    Args:
        config: Configuration dictionary
        
    Returns:
        True if valid, False otherwise
    """
    required_params = ['destination_server', 'destination_port', 'destination_aet', 
                      'local_aet', 'input_directory']
    
    missing_params = []
    for param in required_params:
        if param not in config or not config[param]:
            missing_params.append(param)
    
    if missing_params:
        print(f"Error: Missing required configuration parameters: {', '.join(missing_params)}")
        return False
    
    # Validate port number
    if not isinstance(config['destination_port'], int) or not (1 <= config['destination_port'] <= 65535):
        print(f"Error: Invalid port number: {config['destination_port']}")
        return False
    
    return True


def main():
    """Main function."""
    args = parse_arguments()
    
    # Load configuration
    if args.config:
        config = load_config_file(args.config)
    else:
        # Build config from command line arguments
        config = {
            'destination_server': args.server,
            'destination_port': args.port,
            'destination_aet': args.remote_aet,
            'local_aet': args.local_aet,
            'input_directory': args.input,
            'log_level': args.log_level
        }
    
    # Add num_workers if provided via command line
    if hasattr(args, 'num_workers'):
        config['num_workers'] = args.num_workers
    
    # Validate configuration
    if not validate_config(config):
        sys.exit(1)
    
    # Create transfer instance
    transfer = DicomTransfer(config)
    
    try:
        # Dry run mode - collect files first
        if args.dry_run:
            # Create a temporary finder to count files
            finder_thread = threading.Thread(target=transfer.find_files, args=(config['input_directory'], True))
            finder_thread.start()
            finder_thread.join()
            
            file_list = []
            while not transfer.found_files.empty():
                file_list.append(transfer.found_files.get())
            
            print(f"Dry run: Would transfer {len(file_list)} files:")
            for file_path in file_list[:20]:  # Show first 20
                print(f"  - {file_path}")
            if len(file_list) > 20:
                print(f"  ... and {len(file_list) - 20} more files")
            return
        
        # Start concurrent transfer
        transfer.start_transfer(config['input_directory'])
        
        # Print statistics
        transfer.print_statistics()
        
    except KeyboardInterrupt:
        transfer.logger.info("Interrupted by user, shutting down...")
        transfer.stop()
    except Exception as e:
        transfer.logger.error(f"Fatal error: {e}")
        transfer.stop()
        sys.exit(1)


if __name__ == "__main__":
    main()