flk
465 lines
14 KiB
Python
465 lines
14 KiB
Python
"""
|
||
工具函数模块
|
||
包含: LabelMe 标注解析、linestrip 合并为多边形、多边形转 mask、可视化等
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import math
|
||
import random
|
||
from typing import List, Dict, Tuple, Optional
|
||
|
||
import numpy as np
|
||
from PIL import Image, ImageDraw
|
||
|
||
|
||
# ============================================================
|
||
# 1. 标签映射
|
||
# ============================================================
|
||
|
||
def get_label_map(dataset_dir: str) -> Dict[str, int]:
|
||
"""
|
||
扫描数据集中所有 LabelMe JSON 文件，自动构建标签名→类别ID 映射。
|
||
类别从 1 开始（0 保留给背景）。
|
||
|
||
Args:
|
||
dataset_dir: 数据集目录路径
|
||
|
||
Returns:
|
||
标签名 → 类别ID 的字典，如 {"crosswalk": 1}
|
||
"""
|
||
labels = set()
|
||
for fname in sorted(os.listdir(dataset_dir)):
|
||
if not fname.endswith(".json"):
|
||
continue
|
||
with open(os.path.join(dataset_dir, fname), "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
for shape in data.get("shapes", []):
|
||
labels.add(shape["label"])
|
||
|
||
# 排序保证可复现
|
||
label_map = {label: idx + 1 for idx, label in enumerate(sorted(labels))}
|
||
return label_map
|
||
|
||
|
||
# ============================================================
|
||
# 2. Linestrip 合并为多边形
|
||
# ============================================================
|
||
|
||
def _point_dist(p1, p2) -> float:
|
||
"""计算两点之间的欧氏距离"""
|
||
return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
|
||
|
||
|
||
def merge_linestrips_to_polygon(strips: List[List[List[float]]]) -> List[List[float]]:
|
||
"""
|
||
将多条 linestrip 合并为一个闭合多边形。
|
||
|
||
策略：以第一条 strip 为起点，每次找到与当前多边形末端点距离最近的
|
||
下一条 strip 的端点，将该 strip（可能反转）拼接到多边形末尾，
|
||
最后闭合多边形。
|
||
|
||
对于单条 linestrip，直接闭合。
|
||
|
||
Args:
|
||
strips: 多条 linestrip 的点列表，每条 strip 是 [[x, y], ...]
|
||
|
||
Returns:
|
||
闭合多边形的点列表 [[x, y], ..., [x, y]]（首尾相同）
|
||
"""
|
||
if len(strips) == 0:
|
||
return []
|
||
|
||
# 深拷贝避免修改原始数据
|
||
strips = [list(s) for s in strips]
|
||
|
||
if len(strips) == 1:
|
||
polygon = list(strips[0])
|
||
polygon.append(list(polygon[0])) # 闭合
|
||
return polygon
|
||
|
||
# 从第一条 strip 开始
|
||
polygon = list(strips[0])
|
||
remaining = strips[1:]
|
||
|
||
while remaining:
|
||
last_point = polygon[-1]
|
||
best_idx = 0
|
||
best_dist = float("inf")
|
||
best_reverse = False
|
||
|
||
for i, strip in enumerate(remaining):
|
||
# 尝试正向连接（strip 的起点与 polygon 末端最近）
|
||
d_start = _point_dist(last_point, strip[0])
|
||
# 尝试反向连接（strip 的终点与 polygon 末端最近）
|
||
d_end = _point_dist(last_point, strip[-1])
|
||
|
||
if d_start < best_dist:
|
||
best_dist = d_start
|
||
best_idx = i
|
||
best_reverse = False
|
||
if d_end < best_dist:
|
||
best_dist = d_end
|
||
best_idx = i
|
||
best_reverse = True
|
||
|
||
strip = remaining.pop(best_idx)
|
||
if best_reverse:
|
||
polygon.extend(list(reversed(strip)))
|
||
else:
|
||
polygon.extend(strip)
|
||
|
||
# 闭合多边形
|
||
polygon.append(list(polygon[0]))
|
||
return polygon
|
||
|
||
|
||
# ============================================================
|
||
# 3. 多边形 / Linestrip → Mask
|
||
# ============================================================
|
||
|
||
def shapes_to_masks(
|
||
shapes: List[dict],
|
||
height: int,
|
||
width: int,
|
||
label_map: Dict[str, int],
|
||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||
"""
|
||
将 LabelMe 的 shapes 转换为实例分割所需的 masks、labels、boxes。
|
||
|
||
处理逻辑：
|
||
- polygon 类型：直接作为多边形填充
|
||
- linestrip 类型：将相同 label 的 linestrip 合并为一个多边形后填充
|
||
- rectangle / circle 类型：转换为多边形后处理
|
||
|
||
Args:
|
||
shapes: LabelMe JSON 中的 shapes 列表
|
||
height: 图像高度
|
||
width: 图像宽度
|
||
label_map: 标签名→类别ID 映射
|
||
|
||
Returns:
|
||
masks: [N, H, W] uint8 二值 mask（0 或 1）
|
||
labels: [N] int64 类别ID
|
||
boxes: [N, 4] float32 格式 (x1, y1, x2, y2)
|
||
"""
|
||
# 按 label 分组 linestrip（同 label 的 linestrip 合并为一个实例）
|
||
linestrip_groups: Dict[str, List[List]] = {}
|
||
# 非 linestrip 的 shape 逐个处理
|
||
other_shapes: List[dict] = []
|
||
|
||
for shape in shapes:
|
||
shape_type = shape.get("shape_type", "polygon")
|
||
if shape_type == "linestrip":
|
||
label = shape["label"]
|
||
linestrip_groups.setdefault(label, []).append(shape["points"])
|
||
else:
|
||
other_shapes.append(shape)
|
||
|
||
all_masks = []
|
||
all_labels = []
|
||
|
||
# 处理合并后的 linestrip
|
||
for label, strips in linestrip_groups.items():
|
||
polygon = merge_linestrips_to_polygon(strips)
|
||
if len(polygon) < 3:
|
||
continue
|
||
mask = _polygon_to_mask(polygon, height, width)
|
||
if mask.sum() == 0:
|
||
continue
|
||
all_masks.append(mask)
|
||
all_labels.append(label_map.get(label, 1))
|
||
|
||
# 处理其他 shape 类型
|
||
for shape in other_shapes:
|
||
label = shape["label"]
|
||
shape_type = shape.get("shape_type", "polygon")
|
||
points = shape["points"]
|
||
|
||
if shape_type == "polygon":
|
||
polygon = points
|
||
elif shape_type == "rectangle":
|
||
# 矩形转多边形
|
||
x1, y1 = points[0]
|
||
x2, y2 = points[1]
|
||
polygon = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
|
||
elif shape_type == "circle":
|
||
# 圆形转多边形（近似）
|
||
cx, cy = points[0]
|
||
rx = points[1][0] - cx
|
||
ry = points[1][1] - cy
|
||
polygon = []
|
||
for angle in range(0, 360, 5):
|
||
rad = math.radians(angle)
|
||
polygon.append([cx + rx * math.cos(rad), cy + ry * math.sin(rad)])
|
||
else:
|
||
continue
|
||
|
||
if len(polygon) < 3:
|
||
continue
|
||
|
||
mask = _polygon_to_mask(polygon, height, width)
|
||
if mask.sum() == 0:
|
||
continue
|
||
all_masks.append(mask)
|
||
all_labels.append(label_map.get(label, 1))
|
||
|
||
if len(all_masks) == 0:
|
||
return (
|
||
np.zeros((0, height, width), dtype=np.uint8),
|
||
np.zeros((0,), dtype=np.int64),
|
||
np.zeros((0, 4), dtype=np.float32),
|
||
)
|
||
|
||
masks = np.stack(all_masks, axis=0) # [N, H, W]
|
||
labels = np.array(all_labels, dtype=np.int64)
|
||
|
||
# 从 mask 计算包围框
|
||
boxes = masks_to_boxes(masks)
|
||
|
||
return masks, labels, boxes
|
||
|
||
|
||
def _polygon_to_mask(
|
||
polygon: List[List[float]],
|
||
height: int,
|
||
width: int,
|
||
) -> np.ndarray:
|
||
"""
|
||
将多边形转换为二值 mask。
|
||
|
||
Args:
|
||
polygon: 多边形点列表 [[x, y], ...]
|
||
height: 图像高度
|
||
width: 图像宽度
|
||
|
||
Returns:
|
||
[H, W] uint8 二值 mask（0 或 1）
|
||
"""
|
||
img = Image.new("L", (width, height), 0)
|
||
ImageDraw.Draw(img).polygon(
|
||
[(float(p[0]), float(p[1])) for p in polygon],
|
||
fill=1,
|
||
outline=1,
|
||
)
|
||
return np.array(img, dtype=np.uint8)
|
||
|
||
|
||
def masks_to_boxes(masks: np.ndarray) -> np.ndarray:
|
||
"""
|
||
从一组二值 mask 计算每个实例的包围框。
|
||
|
||
Args:
|
||
masks: [N, H, W] uint8
|
||
|
||
Returns:
|
||
[N, 4] float32, 格式 (x1, y1, x2, y2)
|
||
"""
|
||
boxes = np.zeros((masks.shape[0], 4), dtype=np.float32)
|
||
for i, mask in enumerate(masks):
|
||
ys, xs = np.where(mask > 0)
|
||
if len(xs) == 0:
|
||
continue
|
||
boxes[i] = [xs.min(), ys.min(), xs.max(), ys.max()]
|
||
return boxes
|
||
|
||
|
||
# ============================================================
|
||
# 4. 可视化
|
||
# ============================================================
|
||
|
||
def visualize_instances(
|
||
image: np.ndarray,
|
||
boxes: np.ndarray,
|
||
masks: np.ndarray,
|
||
labels: np.ndarray,
|
||
scores: Optional[np.ndarray] = None,
|
||
label_map: Optional[Dict[int, str]] = None,
|
||
score_threshold: float = 0.5,
|
||
) -> np.ndarray:
|
||
"""
|
||
在图像上绘制实例分割结果（包围框 + mask 叠加）。
|
||
|
||
Args:
|
||
image: [H, W, 3] uint8 BGR 或 RGB 图像
|
||
boxes: [N, 4] 包围框
|
||
masks: [N, H, W] 二值 mask
|
||
labels: [N] 类别ID
|
||
scores: [N] 置信度（可选）
|
||
label_map: 类别ID→名称映射（可选）
|
||
score_threshold: 置信度阈值
|
||
|
||
Returns:
|
||
绘制后的图像 [H, W, 3] uint8
|
||
"""
|
||
try:
|
||
import cv2
|
||
except ImportError:
|
||
# 如果没有 cv2，退化为 PIL 绘制
|
||
return _visualize_instances_pil(
|
||
image, boxes, masks, labels, scores, label_map, score_threshold
|
||
)
|
||
|
||
img = image.copy()
|
||
if img.dtype != np.uint8:
|
||
img = (img * 255).astype(np.uint8)
|
||
|
||
# 生成固定颜色
|
||
np.random.seed(42)
|
||
num_classes = max(int(labels.max()) + 1, 10) if len(labels) > 0 else 10
|
||
colors = np.random.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
|
||
|
||
for i in range(len(boxes)):
|
||
if scores is not None and scores[i] < score_threshold:
|
||
continue
|
||
|
||
color = colors[int(labels[i])].tolist()
|
||
x1, y1, x2, y2 = boxes[i].astype(int)
|
||
|
||
# 绘制 mask 半透明叠加
|
||
mask = masks[i]
|
||
colored_mask = np.zeros_like(img)
|
||
for c in range(3):
|
||
colored_mask[:, :, c] = mask * color[c]
|
||
img = cv2.addWeighted(img, 1.0, colored_mask, 0.5, 0)
|
||
|
||
# 绘制包围框
|
||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
||
|
||
# 绘制标签
|
||
label_text = label_map.get(int(labels[i]), str(labels[i])) if label_map else str(labels[i])
|
||
if scores is not None:
|
||
label_text = f"{label_text}: {scores[i]:.2f}"
|
||
|
||
(tw, th), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
|
||
cv2.rectangle(img, (x1, y1 - th - 6), (x1 + tw + 4, y1), color, -1)
|
||
cv2.putText(
|
||
img, label_text, (x1 + 2, y1 - 4),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA,
|
||
)
|
||
|
||
return img
|
||
|
||
|
||
def _visualize_instances_pil(
|
||
image: np.ndarray,
|
||
boxes: np.ndarray,
|
||
masks: np.ndarray,
|
||
labels: np.ndarray,
|
||
scores: Optional[np.ndarray],
|
||
label_map: Optional[Dict[int, str]],
|
||
score_threshold: float,
|
||
) -> np.ndarray:
|
||
"""使用 PIL 的可视化后备方案（当 cv2 不可用时）"""
|
||
from PIL import Image as PILImage, ImageDraw as PILDraw, ImageFont
|
||
|
||
if image.dtype != np.uint8:
|
||
image = (image * 255).astype(np.uint8)
|
||
|
||
# 转 RGB
|
||
if image.shape[2] == 3:
|
||
pil_img = PILImage.fromarray(image[..., ::-1]) if image.dtype == np.uint8 else PILImage.fromarray(image)
|
||
else:
|
||
pil_img = PILImage.fromarray(image)
|
||
|
||
np.random.seed(42)
|
||
num_classes = max(int(labels.max()) + 1, 10) if len(labels) > 0 else 10
|
||
colors = np.random.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
|
||
|
||
overlay = pil_img.copy()
|
||
draw_overlay = PILDraw.Draw(overlay)
|
||
draw = PILDraw.Draw(pil_img)
|
||
|
||
for i in range(len(boxes)):
|
||
if scores is not None and scores[i] < score_threshold:
|
||
continue
|
||
|
||
color = tuple(colors[int(labels[i])].tolist())
|
||
mask = masks[i]
|
||
# 用颜色填充 mask 区域
|
||
mask_img = PILImage.new("RGB", pil_img.size, (0, 0, 0))
|
||
mask_draw = PILDraw.Draw(mask_img)
|
||
mask_pil = PILImage.fromarray((mask * 255).astype(np.uint8))
|
||
# 叠加半透明 mask
|
||
for c in range(3):
|
||
channel = np.array(pil_img)[:, :, c].astype(float)
|
||
channel[mask > 0] = channel[mask > 0] * 0.5 + color[c] * 0.5
|
||
pil_img_array = np.array(pil_img)
|
||
pil_img_array[:, :, c] = channel
|
||
pil_img = PILImage.fromarray(pil_img_array.astype(np.uint8))
|
||
|
||
draw = PILDraw.Draw(pil_img)
|
||
for i in range(len(boxes)):
|
||
if scores is not None and scores[i] < score_threshold:
|
||
continue
|
||
color = tuple(colors[int(labels[i])].tolist())
|
||
x1, y1, x2, y2 = boxes[i].astype(int)
|
||
draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
|
||
|
||
label_text = label_map.get(int(labels[i]), str(labels[i])) if label_map else str(labels[i])
|
||
if scores is not None:
|
||
label_text = f"{label_text}: {scores[i]:.2f}"
|
||
draw.text((x1 + 2, max(0, y1 - 15)), label_text, fill=color)
|
||
|
||
return np.array(pil_img)
|
||
|
||
|
||
# ============================================================
|
||
# 5. 数据集划分
|
||
# ============================================================
|
||
|
||
def get_image_json_pairs(dataset_dir: str) -> List[Tuple[str, str]]:
|
||
"""
|
||
扫描数据集目录，返回 (image_path, json_path) 配对列表。
|
||
|
||
Args:
|
||
dataset_dir: 数据集目录
|
||
|
||
Returns:
|
||
[(image_path, json_path), ...]
|
||
"""
|
||
pairs = []
|
||
for fname in sorted(os.listdir(dataset_dir)):
|
||
if not fname.endswith(".json"):
|
||
continue
|
||
json_path = os.path.join(dataset_dir, fname)
|
||
with open(json_path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
image_name = data.get("imagePath", fname.replace(".json", ".jpg"))
|
||
image_path = os.path.join(dataset_dir, image_name)
|
||
if not os.path.exists(image_path):
|
||
# 尝试其他常见扩展名
|
||
for ext in [".jpg", ".jpeg", ".png", ".bmp"]:
|
||
alt_path = os.path.join(dataset_dir, fname.replace(".json", ext))
|
||
if os.path.exists(alt_path):
|
||
image_path = alt_path
|
||
break
|
||
pairs.append((image_path, json_path))
|
||
return pairs
|
||
|
||
|
||
def split_dataset(
|
||
pairs: List[Tuple[str, str]],
|
||
val_ratio: float = 0.2,
|
||
seed: int = 42,
|
||
) -> Tuple[List, List]:
|
||
"""
|
||
将数据集划分为训练集和验证集。
|
||
|
||
Args:
|
||
pairs: (image_path, json_path) 配对列表
|
||
val_ratio: 验证集比例
|
||
seed: 随机种子
|
||
|
||
Returns:
|
||
train_pairs, val_pairs
|
||
"""
|
||
rng = random.Random(seed)
|
||
indices = list(range(len(pairs)))
|
||
rng.shuffle(indices)
|
||
val_size = max(1, int(len(pairs) * val_ratio)) if len(pairs) > 1 else 0
|
||
val_indices = set(indices[:val_size])
|
||
train_pairs = [pairs[i] for i in indices if i not in val_indices]
|
||
val_pairs = [pairs[i] for i in val_indices]
|
||
return train_pairs, val_pairs
|