Spaces:
Runtime error
Runtime error
| """PyTorch Dataset class for the APTOS 2019 Blindness Detection challenge.""" | |
| import os | |
| from typing import Callable, Dict, Optional, Tuple, Union | |
| import cv2 | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from torch.utils.data import Dataset | |
| from src.data.preprocessing import ben_graham_pipeline | |
| class APTOSDataset(Dataset): | |
| """APTOS 2019 Diabetic Retinopathy dataset.""" | |
| def __init__( | |
| self, | |
| df_or_csv: Union[pd.DataFrame, str], | |
| images_dir: str, | |
| transform: Optional[Callable] = None, | |
| image_size: int = 380, | |
| apply_ben_graham: bool = True, | |
| is_preprocessed: bool = False, | |
| ) -> None: | |
| """Initialize APTOS dataset. | |
| Args: | |
| df_or_csv: DataFrame or path to CSV with columns ['id_code', 'diagnosis']. | |
| images_dir: Path to image directory containing PNG images. | |
| transform: Optional Albumentations transform pipeline. | |
| image_size: Target image dimension for resizing / preprocessing. | |
| apply_ben_graham: Whether to apply Ben Graham preprocessing on the fly. | |
| is_preprocessed: If True, images in images_dir are already preprocessed. | |
| """ | |
| if isinstance(df_or_csv, str): | |
| self.df = pd.read_csv(df_or_csv) | |
| else: | |
| self.df = df_or_csv.reset_index(drop=True) | |
| self.images_dir = images_dir | |
| self.transform = transform | |
| self.image_size = image_size | |
| self.apply_ben_graham = apply_ben_graham and not is_preprocessed | |
| self.is_preprocessed = is_preprocessed | |
| # Standardize column names | |
| if "id_code" in self.df.columns: | |
| self.image_ids = self.df["id_code"].values | |
| elif "image_id" in self.df.columns: | |
| self.image_ids = self.df["image_id"].values | |
| else: | |
| self.image_ids = self.df.iloc[:, 0].values | |
| if "diagnosis" in self.df.columns: | |
| self.labels = self.df["diagnosis"].values.astype(np.int64) | |
| else: | |
| self.labels = self.df.iloc[:, 1].values.astype(np.int64) | |
| def __len__(self) -> int: | |
| return len(self.df) | |
| def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int, str]: | |
| """Fetch image tensor, integer label, and image path. | |
| Returns: | |
| Tuple of (image_tensor, label, image_path). | |
| """ | |
| img_id = str(self.image_ids[idx]) | |
| # Support both raw id and id with extension | |
| if not img_id.endswith((".png", ".jpg", ".jpeg")): | |
| img_filename = f"{img_id}.png" | |
| else: | |
| img_filename = img_id | |
| img_path = os.path.join(self.images_dir, img_filename) | |
| # Fallback to jpg if png not found | |
| if not os.path.exists(img_path): | |
| alt_path = os.path.join(self.images_dir, f"{img_id}.jpg") | |
| if os.path.exists(alt_path): | |
| img_path = alt_path | |
| image = cv2.imread(img_path) | |
| if image is None: | |
| raise FileNotFoundError(f"Failed to load image at {img_path}") | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| if self.apply_ben_graham: | |
| image = ben_graham_pipeline(image, image_size=self.image_size) | |
| if self.transform is not None: | |
| augmented = self.transform(image=image) | |
| image_tensor = augmented["image"] | |
| else: | |
| # Fallback simple tensor conversion if no transform provided | |
| image_tensor = torch.from_numpy(image.transpose(2, 0, 1)).float() / 255.0 | |
| label = int(self.labels[idx]) | |
| return image_tensor, label, img_path | |
| def get_class_counts(self) -> Dict[int, int]: | |
| """Compute frequency per severity grade.""" | |
| unique, counts = np.unique(self.labels, return_counts=True) | |
| return dict(zip(unique.tolist(), counts.tolist())) | |
| def get_class_weights(self) -> torch.Tensor: | |
| """Compute balanced inverse-frequency class weights for loss functions. | |
| weight_c = total_samples / (num_classes * count_c) | |
| """ | |
| counts = self.get_class_counts() | |
| num_classes = 5 | |
| total = len(self.labels) | |
| weights = [] | |
| for c in range(num_classes): | |
| c_count = counts.get(c, 1) | |
| weights.append(total / (num_classes * c_count)) | |
| return torch.tensor(weights, dtype=torch.float32) | |