Обнавлена структура проекта

datasets/__init__.py

@@ -0,0 +1,7 @@
+from .imagenet import ImageNetDataModule, ImageNetValDataset
+
+
+__all__ = [
+    "ImageNetDataModule",
+    "ImageNetValDataset",
+]

datasets/imagenet.py

@@ -0,0 +1,158 @@
+import torch
+import pandas as pd
+import pytorch_lightning as pl
+
+from PIL import Image
+from pathlib import Path
+from torchvision import datasets
+from torchvision.io import read_file, decode_image, ImageReadMode
+from typing import Optional, Callable, Any
+from torch.utils.data import DataLoader, Dataset
+
+
+TransformType = Optional[Callable[..., Any]]
+
+
+def image_loader(path: str) -> torch.Tensor:
+    data = read_file(path)
+
+    return decode_image(
+        data,
+        mode=ImageReadMode.RGB
+    )
+
+
+class ImageNetValDataset(Dataset):
+    def __init__(
+        self,
+        image_dir,
+        labels_file,
+        class_to_idx,
+        classes,
+        transform=None
+    ):
+        self.image_dir = Path(image_dir)
+        self.transform = transform
+
+        self.class_to_idx = class_to_idx
+        self.classes = classes
+
+        df = pd.read_csv(
+            labels_file,
+            usecols=["ImageId", "PredictionString"]
+        )
+
+        image_ids = df["ImageId"].to_numpy()
+
+        synsets = (
+            df["PredictionString"]
+            .str.split()
+            .str[0]
+            .to_numpy()
+        )
+
+        self.samples = [
+            (
+                self.image_dir / f"{image_id}.JPEG",
+                class_to_idx[synset]
+            )
+            for image_id, synset in zip(image_ids, synsets)
+        ]
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        image_path, label = self.samples[idx]
+
+        image = decode_image(
+            str(image_path),
+            mode=ImageReadMode.RGB
+        )
+
+        if self.transform is not None:
+            image = self.transform(image)
+
+        return image, label
+
+
+class ImageNetDataModule(pl.LightningDataModule):
+    def __init__(self, 
+                 dataset_root: Path,
+                 train_transform: TransformType,
+                 eval_transform: TransformType,
+                 batch_size: int = 64,
+                 num_workers:int = 8
+                 ):
+        super().__init__()
+        
+        self.dataset_root = Path(dataset_root)
+        
+        self.train_transform = train_transform
+        self.eval_transform = eval_transform
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        
+        self.imagenet_root = Path(dataset_root) / "ILSVRC"
+        
+        self.train_dir = self.imagenet_root / "Data" / "CLS-LOC" / "train"
+        self.val_dir = self.imagenet_root / "Data" / "CLS-LOC" / "val"
+        
+        self.mapping_file = Path(dataset_root) / "LOC_synset_mapping.txt"
+        self.val_labels_file = Path(dataset_root) / "LOC_val_solution.csv"
+
+        self.train_dataset = None
+        self.val_dataset = None
+        
+        self.class_to_idx = None
+        self.classes = None
+        
+    def setup(self, stage=None):
+        synset_to_name = {}
+        with open(self.mapping_file, "r", encoding="utf-8") as f:
+            for line in f:
+                synset, class_name = line.strip().split(" ", 1)
+                synset_to_name[synset] = class_name
+        
+        train_dataset = datasets.ImageFolder(root=self.train_dir, 
+                                             transform=self.train_transform,
+                                             loader=image_loader)
+    
+        self.class_to_idx = train_dataset.class_to_idx.copy()
+        self.classes = [synset_to_name[synset] 
+                        for synset in train_dataset.classes]
+        train_dataset.classes = self.classes
+        self.train_dataset = train_dataset
+        
+        self.val_dataset = ImageNetValDataset(
+            image_dir=self.val_dir,
+            labels_file=self.val_labels_file,
+            class_to_idx=self.class_to_idx,
+            classes=self.classes,
+            transform=self.eval_transform
+        )
+        
+    def train_dataloader(self):
+        return DataLoader(
+            dataset=self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            persistent_workers=self.num_workers > 0,
+            prefetch_factor=2 if self.num_workers > 0 else None,
+            drop_last=True
+        )
+    
+    def val_dataloader(self):
+        return DataLoader(
+            dataset=self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            persistent_workers=self.num_workers > 0,
+            prefetch_factor=2 if self.num_workers > 0 else None,
+            drop_last=False
+        )
+