File size: 1,881 Bytes
b5f33fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
from typing import Tuple
import torch
from mivolo.model.mi_volo import MiVOLO
from .age_gender_dataset import AgeGenderDataset
from .age_gender_loader import create_loader
from .classification_dataset import AdienceDataset, FairFaceDataset
DATASET_CLASS_MAP = {
"utk": AgeGenderDataset,
"lagenda": AgeGenderDataset,
"imdb": AgeGenderDataset,
"adience": AdienceDataset,
"fairface": FairFaceDataset,
}
def build(
name: str,
images_path: str,
annotations_path: str,
split: str,
mivolo_model: MiVOLO,
workers: int,
batch_size: int,
) -> Tuple[torch.utils.data.Dataset, torch.utils.data.DataLoader]:
dataset_class = DATASET_CLASS_MAP[name]
dataset: torch.utils.data.Dataset = dataset_class(
images_path=images_path,
annotations_path=annotations_path,
name=name,
split=split,
target_size=mivolo_model.input_size,
max_age=mivolo_model.meta.max_age,
min_age=mivolo_model.meta.min_age,
model_with_persons=mivolo_model.meta.with_persons_model,
use_persons=mivolo_model.meta.use_persons,
disable_faces=mivolo_model.meta.disable_faces,
only_age=mivolo_model.meta.only_age,
)
data_config = mivolo_model.data_config
in_chans = 3 if not mivolo_model.meta.with_persons_model else 6
input_size = (in_chans, mivolo_model.input_size, mivolo_model.input_size)
dataset_loader: torch.utils.data.DataLoader = create_loader(
dataset,
input_size=input_size,
batch_size=batch_size,
mean=data_config["mean"],
std=data_config["std"],
num_workers=workers,
crop_pct=data_config["crop_pct"],
crop_mode=data_config["crop_mode"],
pin_memory=False,
device=mivolo_model.device,
target_type=dataset.target_dtype,
)
return dataset, dataset_loader
|