Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import numpy as np | |
| TARGET_SIZE = (224, 224) | |
| # Read DICOM and return raw pixel array | |
| def dicom_to_pixels(dcm_path: str, save_npz: bool = False, | |
| output_dir: str = None) -> np.ndarray: | |
| import pydicom | |
| ds = pydicom.dcmread(str(dcm_path)) | |
| pixels = ds.pixel_array # matches notebook: ds.pixel_array | |
| if save_npz: | |
| if output_dir is None: | |
| raise ValueError('output_dir must be provided when save_npz=True') | |
| out_path = Path(output_dir) / Path(dcm_path).stem | |
| np.savez_compressed(str(out_path), pixels=pixels) | |
| return pixels | |
| # Load pixels from an npz file | |
| def npz_to_pixels(npz_path): | |
| return np.load(str(npz_path))['pixels'] | |
| # Convert raw pixels to a normalised grayscale 224×224 float32 array including handling accidental images with 3 channels | |
| def pixels_to_gray_resized(pixels): | |
| if pixels.ndim == 3: | |
| pixels = cv2.cvtColor(pixels.astype(np.uint8), cv2.COLOR_BGR2GRAY) | |
| pixels = cv2.resize(pixels.astype(np.float32), TARGET_SIZE) | |
| return (pixels / 255.0).astype('float32') | |
| # Stack the grayscale in three channels for input to the pre-trained model and add batch dimensions | |
| def gray_to_model_input(gray): | |
| rgb = np.stack([gray, gray, gray], axis=-1) | |
| return np.expand_dims(rgb, axis=0).astype('float32') | |
| # Define the full pipeline | |
| def preprocess(source, save_npz=False, output_dir=None): | |
| source = str(source) if not isinstance(source, np.ndarray) else source | |
| if isinstance(source, np.ndarray): | |
| pixels = source.copy() | |
| elif source.endswith('.dcm'): | |
| pixels = dicom_to_pixels(source, save_npz=save_npz, output_dir=output_dir) | |
| elif source.endswith('.npz'): | |
| pixels = npz_to_pixels(source) | |
| else: | |
| pixels = cv2.imread(source, cv2.IMREAD_GRAYSCALE) | |
| if pixels is None: | |
| raise ValueError(f'Could not read image: {source}') | |
| gray = pixels_to_gray_resized(pixels) | |
| img_display = np.clip(gray * 255, 0, 255).astype(np.uint8) | |
| img_input = gray_to_model_input(gray) | |
| return img_display, img_input | |