feat(datasets): YOLO 风格标注数据集与几何/像素增强
- 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>
This commit is contained in:
co-authored by
Claude Opus 5
parent
3a36aad6c3
commit
af22dc639c
@@ -0,0 +1,39 @@
|
||||
"""按配置构建 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,
|
||||
)
|
||||
Reference in New Issue
Block a user