Shaun ZhouandClaude Opus 5
af22dc639c
- yolo: images/ 与同层 labels/ 的路径约定、`class_id cx cy w h` 归一化行解析 （兼容分割多边形取外接矩形）、data.yaml 读写；缺失或空 txt 即无缺陷样本 - film_defect: 抛出图像、框、图像级标签；label_mode 决定多热向量还是互斥下标 - transforms: 图框同步的 Resize / Letterbox / CenterCrop / RandomResizedCrop / 翻转 / ColorJitter / GaussianBlur / PerImageStandardize - target_encoder: 高斯半径与 heatmap / wh / offset / mask 监督信号编码 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
40 lines
1.0 KiB
Python
40 lines
1.0 KiB
Python
"""按配置构建 Dataset / DataLoader。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch.utils.data import DataLoader, Dataset
|
|
|
|
from datasets.film_defect import collate_fn
|
|
from datasets.registry import DATASETS
|
|
|
|
|
|
def build_dataset(cfg: Mapping[str, Any], **kwargs: Any) -> Dataset:
|
|
return DATASETS.build(dict(cfg), **kwargs)
|
|
|
|
|
|
def build_dataloader(
|
|
dataset: Dataset,
|
|
batch_size: int = 8,
|
|
shuffle: bool = False,
|
|
num_workers: int = 4,
|
|
pin_memory: bool | None = None,
|
|
drop_last: bool = False,
|
|
) -> DataLoader:
|
|
"""`pin_memory` 缺省时只在有 CUDA 设备时开启（CPU/MPS 上开启无意义且会告警）。"""
|
|
if pin_memory is None:
|
|
pin_memory = torch.cuda.is_available()
|
|
return DataLoader(
|
|
dataset,
|
|
batch_size=batch_size,
|
|
shuffle=shuffle,
|
|
num_workers=num_workers,
|
|
pin_memory=pin_memory,
|
|
drop_last=drop_last,
|
|
collate_fn=collate_fn,
|
|
persistent_workers=num_workers > 0,
|
|
)
|