159 lines
4.5 KiB
Python
159 lines
4.5 KiB
Python
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
|
|
)
|
|
|