| import cv2 |
| import numpy as np |
| import torch |
| import albumentations as A |
| from albumentations.pytorch import ToTensorV2 |
| from tqdm import tqdm |
| from pathlib import Path |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
|
|
|
|
| class LungSegmenter: |
| """Robust lung field segmentation and text/marker removal using CV techniques.""" |
|
|
| def __init__(self, image_size=224): |
| self.image_size = image_size |
|
|
| def segment_lungs(self, image): |
| """Segment lung fields from chest X-ray. Returns binary mask (0-255 uint8).""" |
| h, w = image.shape |
| if h == 0 or w == 0: |
| return np.ones_like(image, dtype=np.uint8) * 255 |
|
|
| clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8)) |
| enhanced = clahe.apply(image) |
|
|
| blur = cv2.GaussianBlur(enhanced, (7, 7), 0) |
|
|
| _, body_thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) |
| body_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (21, 21)) |
| body_mask = cv2.morphologyEx(body_thresh, cv2.MORPH_CLOSE, body_kernel, iterations=3) |
| body_mask = cv2.morphologyEx(body_mask, cv2.MORPH_OPEN, body_kernel, iterations=1) |
|
|
| inverted = cv2.bitwise_not(blur) |
| masked_inverted = cv2.bitwise_and(inverted, inverted, mask=body_mask) |
| masked_inverted_blur = cv2.GaussianBlur(masked_inverted, (7, 7), 0) |
|
|
| _, lung_thresh = cv2.threshold( |
| masked_inverted_blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU |
| ) |
|
|
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) |
| lung_mask = cv2.morphologyEx(lung_thresh, cv2.MORPH_OPEN, kernel, iterations=2) |
|
|
| contours, _ = cv2.findContours(lung_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
|
|
| min_area = h * w * 0.02 |
| max_area = h * w * 0.45 |
| valid_contours = [] |
| for cnt in contours: |
| area = cv2.contourArea(cnt) |
| if area < min_area or area > max_area: |
| continue |
| x, y, cw, ch = cv2.boundingRect(cnt) |
| if cw < 1 or ch < 1: |
| continue |
| aspect = cw / ch |
| if aspect < 0.3 or aspect > 2.5: |
| continue |
| if y > h * 0.85: |
| continue |
| valid_contours.append(cnt) |
|
|
| if not valid_contours: |
| return np.ones_like(image, dtype=np.uint8) * 255 |
|
|
| mask = np.zeros_like(image, dtype=np.uint8) |
| cv2.drawContours(mask, valid_contours, -1, 255, thickness=cv2.FILLED) |
|
|
| close_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15)) |
| mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, close_kernel, iterations=3) |
|
|
| large_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (31, 31)) |
| mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, large_close, iterations=1) |
|
|
| mask = cv2.GaussianBlur(mask, (15, 15), 0) |
| mask = (mask > 50).astype(np.uint8) * 255 |
|
|
| mask_area_frac = mask.mean() / 255.0 |
| if mask_area_frac < 0.05 or mask_area_frac > 0.75: |
| return np.ones_like(image, dtype=np.uint8) * 255 |
|
|
| return mask |
|
|
| def remove_text_markers(self, image): |
| """Remove text annotations and hardware markers from X-ray.""" |
| h, w = image.shape |
| if h == 0 or w == 0: |
| return image |
|
|
| blur = cv2.GaussianBlur(image, (3, 3), 0) |
|
|
| thresh = cv2.adaptiveThreshold( |
| blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2 |
| ) |
| thresh = cv2.bitwise_not(thresh) |
|
|
| kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2)) |
| thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1) |
|
|
| contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
|
|
| mask = np.zeros_like(image, dtype=np.uint8) |
| for cnt in contours: |
| area = cv2.contourArea(cnt) |
| if area < 5 or area > h * w * 0.01: |
| continue |
| x, y, cw, ch = cv2.boundingRect(cnt) |
| if cw < 1 or ch < 1: |
| continue |
| aspect = cw / max(ch, 1) |
| is_thin_text = aspect > 5 or aspect < 0.2 |
| is_near_edge = ( |
| x < w * 0.03 |
| or x + cw > w * 0.97 |
| or y < h * 0.03 |
| or y + ch > h * 0.97 |
| ) |
| is_circular_marker = 0.8 < aspect < 1.2 and area < h * w * 0.002 |
|
|
| if is_thin_text or is_near_edge or is_circular_marker: |
| cv2.drawContours(mask, [cnt], -1, 255, thickness=cv2.FILLED) |
|
|
| if mask.sum() < 100: |
| return image |
|
|
| result = cv2.inpaint(image, mask, 3, cv2.INPAINT_TELEA) |
| return result |
|
|
|
|
| class LungPreprocessor: |
| def __init__(self, image_size=224): |
| self.image_size = image_size |
| self.segmenter = LungSegmenter(image_size) |
|
|
| def apply_clahe(self, image): |
| clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) |
| return clahe.apply(image) |
|
|
| def normalize_intensity(self, image): |
| p2, p98 = np.percentile(image, (2, 98)) |
| image = np.clip(image, p2, p98) |
| image = ((image - image.min()) / (max(1e-8, image.max() - image.min())) * 255).astype(np.uint8) |
| return image |
|
|
| def preprocess(self, image_path, segment_lung=True): |
| image = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE) |
| if image is None: |
| raise ValueError(f"Failed to load image: {image_path}") |
| if image.mean() < 1 or image.mean() > 253: |
| raise ValueError(f"Image appears corrupted (mean intensity: {image.mean():.1f})") |
|
|
| image = self.segmenter.remove_text_markers(image) |
|
|
| if segment_lung: |
| mask = self.segmenter.segment_lungs(image) |
| mask_float = mask.astype(float) / 255.0 |
| image = (image * mask_float + image * 0.1 * (1 - mask_float)).astype(np.uint8) |
|
|
| image = self.normalize_intensity(image) |
| image = self.apply_clahe(image) |
|
|
| pixel_range = float(image.max()) - float(image.min()) |
| if pixel_range < 2.0: |
| raise ValueError("Preprocessing produced near-uniform image") |
|
|
| return image |
|
|
| def preprocess_array(self, image, segment_lung=True): |
| """Preprocess an in-memory grayscale image array with optional lung segmentation.""" |
| if image is None: |
| raise ValueError("Input image is None") |
| if image.mean() < 1 or image.mean() > 253: |
| raise ValueError(f"Image appears corrupted (mean intensity: {image.mean():.1f})") |
|
|
| image = self.segmenter.remove_text_markers(image) |
|
|
| if segment_lung: |
| mask = self.segmenter.segment_lungs(image) |
| mask_float = mask.astype(float) / 255.0 |
| image = (image * mask_float + image * 0.1 * (1 - mask_float)).astype(np.uint8) |
|
|
| image = self.normalize_intensity(image) |
| image = self.apply_clahe(image) |
|
|
| pixel_range = float(image.max()) - float(image.min()) |
| if pixel_range < 2.0: |
| raise ValueError("Preprocessing produced near-uniform image") |
|
|
| return image |
|
|
| def preprocess_with_mask(self, image_path): |
| """Preprocess image and return (processed_image, lung_mask_at_original_resolution).""" |
| image = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE) |
| if image is None: |
| raise ValueError(f"Failed to load image: {image_path}") |
| if image.mean() < 1 or image.mean() > 253: |
| raise ValueError(f"Image appears corrupted (mean intensity: {image.mean():.1f})") |
|
|
| image = self.segmenter.remove_text_markers(image) |
|
|
| lung_mask = self.segmenter.segment_lungs(image) |
| mask_float = lung_mask.astype(float) / 255.0 |
| image = (image * mask_float + image * 0.1 * (1 - mask_float)).astype(np.uint8) |
|
|
| image = self.normalize_intensity(image) |
| image = self.apply_clahe(image) |
|
|
| pixel_range = float(image.max()) - float(image.min()) |
| if pixel_range < 2.0: |
| raise ValueError("Preprocessing produced near-uniform image") |
|
|
| return image, lung_mask |
|
|
|
|
| def get_train_transforms(image_size=224): |
| return A.Compose([ |
| A.Resize(height=image_size, width=image_size), |
| A.RandomResizedCrop(size=(image_size, image_size), scale=(0.85, 1.0), p=0.5), |
| A.HorizontalFlip(p=0.5), |
| A.Rotate(limit=10, p=0.7), |
| A.RandomBrightnessContrast(brightness_limit=0.15, contrast_limit=0.15, p=0.5), |
| A.GaussNoise(var_limit=(10.0, 50.0), p=0.3), |
| A.GridDistortion(num_steps=5, distort_limit=0.1, p=0.2), |
| A.Normalize(mean=[0.5], std=[0.5], max_pixel_value=255.0), |
| ToTensorV2() |
| ]) |
|
|
|
|
| def get_val_transforms(image_size=224): |
| return A.Compose([ |
| A.Resize(height=image_size, width=image_size), |
| A.Normalize(mean=[0.5], std=[0.5], max_pixel_value=255.0), |
| ToTensorV2() |
| ]) |
|
|
|
|
| class PreprocessedDataset(torch.utils.data.Dataset): |
| def __init__(self, image_paths, labels, transforms=None, use_preprocessing=True, cache_pt=True, load_to_ram=False): |
| self.image_paths = image_paths |
| self.labels = labels |
| self.transforms = transforms |
| self.use_preprocessing = use_preprocessing |
| self.cache_pt = cache_pt |
| self.preprocessor = LungPreprocessor() if use_preprocessing else None |
| self.ram_cache = None |
|
|
| if load_to_ram: |
| self._preload_to_ram() |
|
|
| def _get_preprocessed(self, img_path, target_size=224): |
| if not self.use_preprocessing: |
| image = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE) |
| image = cv2.resize(image, (target_size, target_size)) |
| return image |
|
|
| cache_path = img_path.with_suffix('.pt.cache') |
| if self.cache_pt and cache_path.exists(): |
| return torch.load(cache_path, weights_only=True).numpy() |
|
|
| image = self.preprocessor.preprocess(str(img_path), segment_lung=True) |
|
|
| if self.cache_pt: |
| torch.save(torch.from_numpy(image), cache_path) |
|
|
| return image |
|
|
| def _build_cache(self, paths): |
| processor = LungPreprocessor() |
| for p in paths: |
| try: |
| img = processor.preprocess(str(p), segment_lung=True) |
| torch.save(torch.from_numpy(img), p.with_suffix('.pt.cache')) |
| except Exception: |
| pass |
|
|
| def _preload_to_ram(self): |
| n_total = len(self.image_paths) |
| print(f" Pre-loading {n_total} images into RAM...") |
|
|
| uncached = [p for p in self.image_paths if not p.with_suffix('.pt.cache').exists()] |
| if uncached: |
| n_workers = min(8, len(uncached)) |
| chunk_size = (len(uncached) + n_workers - 1) // n_workers |
| chunks = [uncached[i:i+chunk_size] for i in range(0, len(uncached), chunk_size)] |
| print(f" Processing {len(uncached)} uncached images with {n_workers} workers...") |
| worker = LungPreprocessor() |
| for chunk in tqdm(chunks, desc="Caching"): |
| for p in chunk: |
| try: |
| img = worker.preprocess(str(p), segment_lung=True) |
| torch.save(torch.from_numpy(img), p.with_suffix('.pt.cache')) |
| except Exception: |
| pass |
|
|
| self.ram_cache = [] |
| target_size = None |
| if self.transforms: |
| for t in self.transforms: |
| if hasattr(t, 'height'): |
| target_size = t.height |
| break |
| if target_size is None: |
| target_size = self.preprocessor.image_size if self.preprocessor else 224 |
|
|
| for path in tqdm(self.image_paths, desc="Loading RAM"): |
| try: |
| image = torch.load(path.with_suffix('.pt.cache'), weights_only=True).numpy() |
| self.ram_cache.append(image) |
| except Exception as e: |
| print(f" Warning: Failed to load {path}: {e}") |
| self.ram_cache.append(np.zeros((target_size, target_size), dtype=np.uint8)) |
|
|
| n = len(self.ram_cache) |
| size_gb = n * target_size * target_size * 1 / 1e9 |
| print(f" Cached {n} images in RAM (~{size_gb:.1f} GB)") |
|
|
| def __len__(self): |
| return len(self.image_paths) |
|
|
| def __getitem__(self, idx): |
| label = self.labels[idx] |
|
|
| if self.ram_cache is not None: |
| image = self.ram_cache[idx] |
| elif self.use_preprocessing: |
| image = self._get_preprocessed(self.image_paths[idx]) |
| else: |
| image = cv2.imread(str(self.image_paths[idx]), cv2.IMREAD_GRAYSCALE) |
| target_size = self.transforms[0].height if self.transforms and hasattr(self.transforms[0], 'height') else 224 |
| image = cv2.resize(image, (target_size, target_size)) |
|
|
| if self.transforms: |
| augmented = self.transforms(image=image) |
| image = augmented['image'] |
| else: |
| image = torch.from_numpy(image).unsqueeze(0).float() / 255.0 |
|
|
| if image.shape[0] == 3: |
| image = image.mean(dim=0, keepdim=True) |
| elif image.dim() == 2: |
| image = image.unsqueeze(0) |
|
|
| return image, label |
|
|
