Update preprocessing.py

preprocessing.py

@@ -0,0 +1,105 @@
+"""
+Preprocessing script for DICOM medical images
+"""
+import os
+import numpy as np
+import pydicom
+import cv2
+from glob import glob
+from tqdm import tqdm
+import argparse
+
+def apply_window_level(image, window_center, window_width):
+    """
+    Apply windowing to the DICOM image to enhance visualization
+    
+    Args:
+        image (numpy.ndarray): Input image
+        window_center (float): Window center (level)
+        window_width (float): Window width
+        
+    Returns:
+        numpy.ndarray: Windowed image
+    """
+    img_min = window_center - window_width // 2
+    img_max = window_center + window_width // 2
+    
+    windowed = np.clip(image, img_min, img_max)
+    windowed = (windowed - img_min) / (img_max - img_min) * 255.0
+    
+    return windowed.astype(np.uint8)
+
+def process_dicom_files(input_dir, output_dir):
+    """
+    Process all DICOM files in the input directory
+    
+    Args:
+        input_dir (str): Directory containing DICOM files
+        output_dir (str): Directory to save processed images
+    """
+    # Create output directories
+    os.makedirs(output_dir, exist_ok=True)
+    images_dir = os.path.join(output_dir, "images")
+    os.makedirs(images_dir, exist_ok=True)
+    
+    # Find all DICOM files
+    dicom_files = glob(os.path.join(input_dir, "**", "*.dcm"), recursive=True)
+    print(f"Found {len(dicom_files)} DICOM files")
+    
+    # Process each file
+    for i, dicom_path in enumerate(tqdm(dicom_files, desc="Processing DICOM files")):
+        try:
+            # Read DICOM file
+            dicom = pydicom.dcmread(dicom_path)
+            
+            # Extract image data
+            image = dicom.pixel_array
+            
+            # Apply windowing if available
+            if hasattr(dicom, 'WindowCenter') and hasattr(dicom, 'WindowWidth'):
+                window_center = dicom.WindowCenter
+                window_width = dicom.WindowWidth
+                
+                # Handle multiple window values
+                if isinstance(window_center, pydicom.multival.MultiValue):
+                    window_center = window_center[0]
+                if isinstance(window_width, pydicom.multival.MultiValue):
+                    window_width = window_width[0]
+                    
+                image = apply_window_level(image, window_center, window_width)
+            else:
+                # Apply default windowing for CT images
+                if dicom.Modality == "CT":
+                    image = apply_window_level(image, 40, 400)  # Soft tissue window
+                else:
+                    # Normalize to 0-255 range
+                    image = ((image - image.min()) / (image.max() - image.min() + 1e-8) * 255).astype(np.uint8)
+            
+            # Generate output filename
+            patient_id = dicom.PatientID if hasattr(dicom, 'PatientID') else "unknown"
+            series_uid = dicom.SeriesInstanceUID if hasattr(dicom, 'SeriesInstanceUID') else "unknown"
+            instance_uid = dicom.SOPInstanceUID if hasattr(dicom, 'SOPInstanceUID') else str(i)
+            
+            output_filename = f"{patient_id}_{series_uid[-8:]}_{instance_uid[-8:]}.png"
+            output_path = os.path.join(images_dir, output_filename)
+            
+            # Save image
+            cv2.imwrite(output_path, image)
+            
+            # Save metadata
+            with open(os.path.join(output_dir, "metadata.txt"), "a") as f:
+                f.write(f"{output_filename},{dicom_path}\n")
+                
+        except Exception as e:
+            print(f"Error processing {dicom_path}: {e}")
+    
+    print(f"Processing complete. Processed images saved to {images_dir}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Process DICOM files')
+    parser.add_argument('--input', type=str, required=True, help='Input directory with DICOM files')
+    parser.add_argument('--output', type=str, default='./processed_data', help='Output directory')
+    
+    args = parser.parse_args()
+    
+    process_dicom_files(args.input, args.output)