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([])
|