Mask R-CNN instance segmentation training code
This commit is contained in:
flk
+227
@@ -0,0 +1,227 @@
|
||||
"""
|
||||
自定义 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([])
|
||||
Reference in New Issue
Block a user