flk
228 lines
6.9 KiB
Python
228 lines
6.9 KiB
Python
"""
|
||
自定义 Dataset 类
|
||
支持 LabelMe 标注格式（polygon / linestrip / rectangle / circle）
|
||
内置数据增强（水平翻转、亮度调整）
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
from typing import List, Dict, Tuple, Optional
|
||
|
||
import numpy as np
|
||
from PIL import Image
|
||
import torch
|
||
from torch.utils.data import Dataset
|
||
|
||
from utils import (
|
||
get_label_map,
|
||
shapes_to_masks,
|
||
get_image_json_pairs,
|
||
split_dataset,
|
||
)
|
||
|
||
|
||
class LabelMeDataset(Dataset):
|
||
"""
|
||
LabelMe 格式的实例分割数据集。
|
||
|
||
数据目录结构:
|
||
dataset/
|
||
├── 1.jpg
|
||
├── 1.json
|
||
├── 2.jpg
|
||
├── 2.json
|
||
└── ...
|
||
|
||
每个 JSON 文件包含 LabelMe 标注，shapes 列表中每个 shape 有:
|
||
- label: 类别名
|
||
- points: [[x, y], ...] 点坐标
|
||
- shape_type: "polygon" / "linestrip" / "rectangle" / "circle"
|
||
|
||
linestrip 类型会自动将相同 label 的多条线段合并为一个多边形实例。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
dataset_dir: str,
|
||
label_map: Optional[Dict[str, int]] = None,
|
||
pairs: Optional[List[Tuple[str, str]]] = None,
|
||
transforms: Optional[object] = None,
|
||
):
|
||
"""
|
||
Args:
|
||
dataset_dir: 数据集目录路径
|
||
label_map: 标签名→类别ID 映射（若为 None 则自动构建）
|
||
pairs: 指定使用的 (image_path, json_path) 配对（若为 None 则扫描全部）
|
||
transforms: 数据增强变换
|
||
"""
|
||
self.dataset_dir = dataset_dir
|
||
self.transforms = transforms
|
||
|
||
# 自动构建或使用传入的标签映射
|
||
if label_map is None:
|
||
label_map = get_label_map(dataset_dir)
|
||
self.label_map = label_map
|
||
|
||
# 获取图像-标注配对
|
||
if pairs is None:
|
||
pairs = get_image_json_pairs(dataset_dir)
|
||
self.pairs = pairs
|
||
|
||
def __len__(self) -> int:
|
||
return len(self.pairs)
|
||
|
||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict]:
|
||
image_path, json_path = self.pairs[idx]
|
||
|
||
# 读取图像
|
||
image = Image.open(image_path).convert("RGB")
|
||
width, height = image.size
|
||
|
||
# 读取标注
|
||
with open(json_path, "r", encoding="utf-8") as f:
|
||
anno = json.load(f)
|
||
|
||
# 确保 imageHeight / imageWidth 与实际一致
|
||
anno_h = anno.get("imageHeight", height)
|
||
anno_w = anno.get("imageWidth", width)
|
||
|
||
# 转换 shapes → masks, labels, boxes
|
||
masks, labels, boxes = shapes_to_masks(
|
||
anno.get("shapes", []),
|
||
anno_h,
|
||
anno_w,
|
||
self.label_map,
|
||
)
|
||
|
||
# 转为 numpy 数组
|
||
image_np = np.array(image, dtype=np.float32) / 255.0 # [H, W, 3] 0~1
|
||
|
||
# 构建 target 字典
|
||
num_instances = len(labels)
|
||
target = {
|
||
"boxes": torch.as_tensor(boxes, dtype=torch.float32),
|
||
"labels": torch.as_tensor(labels, dtype=torch.int64),
|
||
"masks": torch.as_tensor(masks, dtype=torch.uint8),
|
||
"image_id": torch.tensor([idx]),
|
||
"area": torch.zeros((num_instances,), dtype=torch.float32),
|
||
"iscrowd": torch.zeros((num_instances,), dtype=torch.int64),
|
||
}
|
||
|
||
# 计算 area
|
||
if num_instances > 0:
|
||
target["area"] = (target["masks"] > 0).sum(dim=(1, 2)).float()
|
||
|
||
# 应用数据增强
|
||
if self.transforms is not None:
|
||
image_np, target = self.transforms(image_np, target)
|
||
|
||
# 转换为 tensor [C, H, W]
|
||
image_tensor = torch.as_tensor(image_np, dtype=torch.float32)
|
||
if image_tensor.dim() == 3:
|
||
image_tensor = image_tensor.permute(2, 0, 1)
|
||
|
||
return image_tensor, target
|
||
|
||
def get_label_map(self) -> Dict[str, int]:
|
||
return self.label_map
|
||
|
||
def get_num_classes(self) -> int:
|
||
"""返回类别数（不含背景）"""
|
||
return max(self.label_map.values()) if self.label_map else 0
|
||
|
||
|
||
# ============================================================
|
||
# 数据增强
|
||
# ============================================================
|
||
|
||
class Compose:
|
||
"""组合多个变换"""
|
||
|
||
def __init__(self, transforms: List):
|
||
self.transforms = transforms
|
||
|
||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||
for t in self.transforms:
|
||
image, target = t(image, target)
|
||
return image, target
|
||
|
||
|
||
class RandomHorizontalFlip:
|
||
"""随机水平翻转（同时翻转图像、mask、boxes）"""
|
||
|
||
def __init__(self, prob: float = 0.5):
|
||
self.prob = prob
|
||
|
||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||
if np.random.random() < self.prob:
|
||
image = np.ascontiguousarray(image[:, ::-1, :])
|
||
_, width = image.shape[:2]
|
||
|
||
# 翻转 mask
|
||
if "masks" in target and len(target["masks"]) > 0:
|
||
target["masks"] = torch.flip(target["masks"], dims=[2])
|
||
target["masks"] = target["masks"].contiguous()
|
||
|
||
# 翻转 boxes
|
||
if "boxes" in target and len(target["boxes"]) > 0:
|
||
boxes = target["boxes"].clone()
|
||
boxes[:, 0] = width - target["boxes"][:, 2]
|
||
boxes[:, 2] = width - target["boxes"][:, 0]
|
||
target["boxes"] = boxes
|
||
|
||
return image, target
|
||
|
||
|
||
class RandomBrightness:
|
||
"""随机亮度调整"""
|
||
|
||
def __init__(self, brightness_range: float = 0.2):
|
||
self.brightness_range = brightness_range
|
||
|
||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||
factor = 1.0 + np.random.uniform(
|
||
-self.brightness_range, self.brightness_range
|
||
)
|
||
image = np.clip(image * factor, 0.0, 1.0)
|
||
return image, target
|
||
|
||
|
||
class RandomContrast:
|
||
"""随机对比度调整"""
|
||
|
||
def __init__(self, contrast_range: float = 0.2):
|
||
self.contrast_range = contrast_range
|
||
|
||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||
factor = 1.0 + np.random.uniform(
|
||
-self.contrast_range, self.contrast_range
|
||
)
|
||
mean = image.mean()
|
||
image = np.clip((image - mean) * factor + mean, 0.0, 1.0)
|
||
return image, target
|
||
|
||
|
||
class Normalize:
|
||
"""标准化（使用 ImageNet 均值和标准差）"""
|
||
|
||
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||
|
||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||
image = (image - self.MEAN) / self.STD
|
||
return image, target
|
||
|
||
|
||
def get_train_transforms() -> Compose:
|
||
"""获取训练集数据增强"""
|
||
return Compose([
|
||
RandomHorizontalFlip(prob=0.5),
|
||
RandomBrightness(brightness_range=0.2),
|
||
RandomContrast(contrast_range=0.2),
|
||
])
|
||
|
||
|
||
def get_val_transforms() -> Compose:
|
||
"""获取验证集数据增强（仅标准化，不做随机变换）"""
|
||
return Compose([])
|