最近是接了一个需求咨询图像处理类的，甲方要在卡车过磅的地方装一个摄像头用检测卡车的车斗雨覆是否完全， 让我大致理了下需求并对技术核心做下预研究
开发一套图像处理软件，能够实时监控经过的卡车并判断其车斗的雨覆状态。
系统需具备以下功能：
- 图像采集：通过高分辨率摄像头采集卡车经过时的图像。
- 图像处理：对采集的图像进行处理，识别车斗及雨覆的具体位置。
- 状态判断：
- 判断雨覆是否完全覆盖车斗。
- 在雨覆未完全覆盖时，生成警报或提示信息。
例如这样的图片，检测上方雨覆是否完全遮盖住
需求分析
图像采集是视频流接收，确保摄像头支持所选协议，并具备高清分辨率（至少1080p）以提高图像识别的准确性。摄像头应具备良好的低光性能，以适应不同的环境光照条件。
- 使用开源媒体框架（如 FFmpeg 或 GStreamer）来接收和处理视频流。
- 实现视频流的解码，提取每一帧图像供后续处理。
所以在技术预研上直接从图片开始， 目测了下需要使用图像语义分割再在分割图像基础上再计算雨覆的遮盖率，会使用到的工具具体有
导入必要的库，用于深度学习（PyTorch）、图像处理（PIL、OpenCV）、可视化（Matplotlib）、以及 YOLOv8（Ultralytics）目标检测和 DINOv2（Transformers）语义分割。
- PyTorch的作用：
- 提供核心深度学习功能，如 torch.nn 用于定义 DINOv2 分割模型（DINOv2ForSegmentation 类）。
- 通过 torch.device 确定设备（CPU 或 CUDA），将模型和数据加载到 GPU（self.device 和 self.model.to(self.device)）。
- 处理张量操作（如 torch.softmax 在 postprocess 中生成概率分布）。
- PIL 作用：用于图像处理，支持打开、转换、调整大小和增强图像。
- 加载和处理图像文件（如 Image.open("trunk.jpg")）。
- 转换图像格式（如 image.convert("RGB")），调整大小（如 mask.resize），并支持数据增强（如 ImageEnhance 进行亮度、对比度、色调和饱和度调整）。
- 创建和保存掩码或结果图像（如 Image.fromarray 和 save 方法）。
- cv2的作用：
- 处理颜色分割（如 color_based_segmentation 使用 HSV 颜色空间分割车斗和覆盖布）。
- 进行后处理优化，包括形态学操作（dilate、erode、morphologyEx 在 enhance_segmentation 和 color_based_segmentation 中填补空洞、去除噪声）。
- 边缘检测（Canny 在 enhance_segmentation 中捕捉车斗边缘）和轮廓检测（findContours 填补完整轮廓）。
- Matplotlib作用：用于数据可视化和绘图，适合生成图形和保存图像。
- 保存分割结果的图像（如 plt.imsave 在 visualize_and_extract_regions 中保存车斗和覆盖布到黑色背景的图片）。
- 提供可视化支持，但当前代码未直接使用 Matplotlib 绘制图表，仅用于文件保存
- Ultralytics（YOLOv8）的作用：
- 加载预训练的 yolov8n.pt 模型（YOLO("yolov8n.pt")），检测图像中的卡车（detect_truck 函数）。
- 返回卡车的边界框（box.xyxy），用于裁剪图像区域供 DINOv2 语义分割，确保仅在 truck 区域内分割车斗和覆盖布。
- 设置置信度阈值（conf=0.3）以平衡检测精度和召回率。
定义类别和参数（Constants and Parameters）
CLASS_NAMES = ["background", "truck_bed", "tarp"]
AUGMENTATION_PARAMS = {
"brightness_factor": (0.8, 1.2),
"contrast_factor": (0.8, 1.2),
"rotation_range": (-30, 30),
"hue_shift": (-0.1, 0.1),
"saturation_factor": (0.8, 1.2),
}
YOLO_CLASSES = {
6: "train",
7: "truck",
}- 作用：
自定义 DINOv2 分割模型（DINOv2ForSegmentation）
class DINOv2ForSegmentation(nn.Module):
def __init__(self, num_classes=3, model_name="./dinov2_base/"):
# 加载 DINOv2 主干网络并冻结参数
self.backbone = ViTModel.from_pretrained(model_name)
hidden_size = self.backbone.config.hidden_size
# 添加分割头，适配 518x518 输入
self.segmentation_head = nn.Sequential(
nn.Conv2d(hidden_size, 256, kernel_size=1),
nn.Upsample(scale_factor=14, mode='bilinear'),
nn.Conv2d(256, num_classes, kernel_size=1)
)
for param in self.backbone.parameters():
param.requires_grad = False
def forward(self, pixel_values):
# 从 DINOv2 提取特征并通过分割头生成分割结果
outputs = self.backbone(pixel_values)
last_hidden = outputs.last_hidden_state
features = last_hidden[:, 1:].permute(0, 2, 1).view(last_hidden.size(0), -1, 37, 37)
logits = self.segmentation_head(features)
return logits作用：
- 定义基于 DINOv2 的语义分割模型，使用预训练的 ViT（Vision Transformer）作为主干网络，冻结其参数以减少计算量。
- 添加自定义分割头（segmentation_head），将 37x37 的特征图上采样并生成 3 类的分割结果（背景、车斗、覆盖布）。
- forward 方法处理输入图像（518x518），输出分割 logits。
分割管道（SegmentationPipeline）
class SegmentationPipeline:
def __init__(self, num_classes=3):
# 初始化 DINOv2 模型和特征提取器
self.feature_extractor = ViTFeatureExtractor.from_pretrained("./dinov2_base", size={"height": 518, "width": 518})
self.model = DINOv2ForSegmentation(num_classes)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
self.model.eval()
def augment_image(self, image):
# 应用数据增强（亮度、对比度、旋转、色调、饱和度）
image = ImageEnhance.Brightness(image).enhance(random.uniform(*AUGMENTATION_PARAMS["brightness_factor"]))
image = ImageEnhance.Contrast(image).enhance(random.uniform(*AUGMENTATION_PARAMS["contrast_factor"]))
image = ImageEnhance.Color(image).enhance(random.uniform(*AUGMENTATION_PARAMS["saturation_factor"]))
angle = random.uniform(*AUGMENTATION_PARAMS["rotation_range"])
image = image.rotate(angle, expand=True, fillcolor=(0, 0, 0))
return image
def preprocess(self, image):
# 预处理图像：转换为 RGB，应用增强，调整到 518x518
if image.mode != "RGB":
image = image.convert("RGB")
augmented_image = self.augment_image(image)
inputs = self.feature_extractor(images=augmented_image, return_tensors="pt", size={"height": 518, "width": 518})
return inputs.pixel_values.to(self.device)
def postprocess(self, logits, original_size):
# 后处理：softmax 转换为概率，取最大值生成掩码，调整回原图大小
probs = torch.softmax(logits, dim=1)
mask = torch.argmax(probs, dim=1).squeeze().cpu().numpy()
mask = Image.fromarray(mask.astype(np.uint8)).resize(original_size, Image.NEAREST)
return mask
def predict(self, image, truck_bbox=None):
# 在 truck 边界框内或全图进行预测
if truck_bbox:
cropped_image = image.crop((
max(0, truck_bbox[0] - 20), max(0, truck_bbox[1] - 20),
min(image.size[0], truck_bbox[2] + 20), min(image.size[1], truck_bbox[3] + 20)
))
else:
cropped_image = image
inputs = self.preprocess(cropped_image)
with torch.no_grad():
logits = self.model(inputs)
probs = torch.softmax(logits, dim=1)
mask = self.postprocess(logits, cropped_image.size)
if truck_bbox:
full_mask = Image.new("L", image.size, 0)
full_mask.paste(mask, (max(0, truck_bbox[0] - 20), max(0, truck_bbox[1] - 20),
min(image.size[0], truck_bbox[2] + 20), min(image.size[1], truck_bbox[3] + 20)))
return full_mask
return mask作用
- SegmentationPipeline 封装了 DINOv2 模型的预处理、预测和后处理逻辑。
- augment_image：通过随机变换增强图像，模拟不同光照、角度和颜色，提高模型泛化能力。
- preprocess：将输入图像转换为 RGB，应用增强，调整到 518x518，发送到 GPU/CPU。
- postprocess：将模型输出转换为掩码，调整回原图大小。
- predict：根据 YOLO 检测的 truck 边界框裁剪图像进行分割，支持扩展边界（padding=20）以捕捉边缘。
可视化和提取区域（visualize_and_extract_regions）
def visualize_and_extract_regions(original_image, mask, save_base_path="output", num_classes=3):
# 创建黑色背景，提取车斗和覆盖布，保存到单独图片
colormap = np.array([[0, 0, 0], [255, 0, 0], [0, 255, 0]])
mask_array = np.array(mask.convert("L")).resize(original_image.size, Image.NEAREST)
black_background = np.zeros((*original_image.size, 3), dtype=np.uint8)
for class_id, class_name in enumerate(CLASS_NAMES[1:], 1):
class_mask = (mask_array == class_id).astype(np.uint8) * 255
extracted_region = black_background.copy()
for i in range(3):
extracted_region[:, :, i] = black_background[:, :, i] * (1 - class_mask / 255) + colormap[class_id, i] * (class_mask / 255)
Image.fromarray(extracted_region).save(f"{save_base_path}_{class_name}_on_black.png")
mask.save(f"{save_base_path}_mask.png")
color_based_mask = color_based_segmentation(original_image)
if color_based_mask:
visualize_color_based_mask_on_black(original_image, color_based_mask, save_base_path + "_color_based_on_black")作用：
- 将车斗和覆盖布提取到黑色背景的单独图片，忽略背景。
- 调整掩码尺寸匹配原图，确保尺寸一致。
- 保存原始掩码和颜色增强后的掩码，用于检查和调试。
YOLOv8 检测卡车（detect_truck）
def detect_truck(image_path):
# 使用 YOLOv8 检测卡车，返回边界框
model = YOLO("yolov8n.pt")
print('yolov8n load successfully~~~')
results = model(image_path, conf=0.3)
for result in results:
for box in result.boxes:
cls = int(box.cls[0])
if YOLO_CLASSES.get(cls):
print('detect out:', YOLO_CLASSES.get(cls))
x1, y1, x2, y2 = box.xyxy[0].tolist()
return (int(x1), int(y1), int(x2), int(y2))
return None- 加载 YOLOv8 模型（yolov8n.pt），检测图像中的卡车（类别 "truck"，ID 7）。
- 设置置信度阈值（conf=0.3），返回卡车的边界框（x1, y1, x2, y2）。
颜色分割（color_based_segmentation）
def color_based_segmentation(image):
# 使用颜色阈值分割车斗（红、蓝、黄）和覆盖布（绿、蓝）
img = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
mask_truck = cv2.bitwise_or(cv2.bitwise_or(cv2.inRange(hsv, [0, 120, 70], [10, 255, 255]),
cv2.inRange(hsv, [170, 120, 70], [180, 255, 255])),
cv2.bitwise_or(cv2.inRange(hsv, [100, 120, 70], [130, 255, 255]),
cv2.inRange(hsv, [20, 120, 70], [40, 255, 255])))
mask_tarp = cv2.bitwise_or(cv2.inRange(hsv, [35, 40, 40], [85, 255, 255]),
cv2.inRange(hsv, [100, 40, 40], [130, 255, 255]))
combined_mask = np.zeros(hsv.shape[:2], dtype=np.uint8)
combined_mask[mask_truck > 0] = 1
combined_mask[mask_tarp > 0] = 2
kernel = np.ones((5, 5), np.uint8)
return Image.fromarray(cv2.morphologyEx(cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel),
cv2.MORPH_OPEN, kernel))-
作用：
- 使用 HSV 颜色空间分割车斗（红、蓝、黄）和覆盖布（绿、蓝），提高对颜色变化的鲁棒性。
- 应用形态学操作（闭运算和开运算）去除噪声并填补空洞，作为 DINOv2 的补充。
def visualize_color_based_mask_on_black(original_image, mask, save_path):
# 将颜色分割结果提取到黑色背景
colormap = np.array([[0, 0, 0], [255, 0, 0], [0, 255, 0]])
mask_array = np.array(mask)
black_background = np.zeros((*mask_array.shape, 3), dtype=np.uint8)
for class_id in [1, 2]:
class_mask = (mask_array == class_id).astype(np.uint8) * 255
for i in range(3):
black_background[:, :, i] = black_background[:, :, i] * (1 - class_mask / 255) + colormap[class_id, i] * (class_mask / 255)
Image.fromarray(black_background).save(save_path + ".png")- 作用：将颜色分割的车斗和覆盖布提取到黑色背景，生成单独的图像文件。
测试卡车图片语义分割
if __name__ == "__main__":
# 初始化 pipeline
num_classes = 3 # 背景, 车斗, 覆盖布
try:
pipeline = SegmentationPipeline(num_classes=num_classes)
# 加载原始图像
image_path = "trunk.jpg"
original_image = Image.open(image_path)
# 使用YOLOv8检测卡车
truck_bbox = detect_truck(image_path)
if truck_bbox:
print(f"Detected truck bounding box: {truck_bbox}")
# 进行分割
segmentation_mask = pipeline.predict(original_image, truck_bbox)
# 提取并可视化车斗和覆盖布到黑色背景
visualize_and_extract_regions(
original_image=original_image,
mask=segmentation_mask,
save_base_path="segmentation_output",
num_classes=num_classes
)
else:
print("No truck detected, using full image for segmentation")
except Exception as e:
print(f"Error in main execution: {e}")运行~~
分割后图像
优化目标和当前问题
- 当前问题：车斗边缘未完全分割，可能因 DINOv2 模型未训练捕捉细小边缘、YOLO 边界框未包含边缘、或后处理未充分扩展轮廓。
- 优化：通过扩展 YOLO 边界框（padding）、增强边缘检测（Canny 和形态学操作）、结合颜色和形状信息，确保车斗完整轮廓被分割。
经过修改和测试
import torch
import torch.nn as nn
from transformers import ViTModel, ViTFeatureExtractor
from PIL import Image, ImageEnhance
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import os
import cv2 # For color-based and shape-based post-processing
import random
from ultralytics import YOLO # For YOLOv8
# 定义类别：0 = 背景, 1 = 车斗, 2 = 覆盖布
CLASS_NAMES = ["background", "truck_bed", "tarp"]
# 数据增强参数
AUGMENTATION_PARAMS = {
"brightness_factor": (0.8, 1.2), # 亮度变化范围
"contrast_factor": (0.8, 1.2), # 对比度变化范围
"rotation_range": (-30, 30), # 旋转角度范围（度）
"hue_shift": (-0.1, 0.1), # 色调变化范围
"saturation_factor": (0.8, 1.2), # 饱和度变化范围
}
# YOLOv8 类别映射（假设 'truck' 是可识别的类别）
YOLO_CLASSES = {
6: "train",
7: "truck",
}
# ======================
# 自定义分割模型定义（适配518x518）
# ======================
class DINOv2ForSegmentation(nn.Module):
def __init__(self, num_classes=3, model_name="./dinov2_base/"):
super().__init__()
# 加载 DINOv2 主干网络
try:
self.backbone = ViTModel.from_pretrained(model_name)
except Exception as e:
print(f"Error loading DINOv2 model: {e}")
raise
hidden_size = self.backbone.config.hidden_size
# 分割头调整（适配518输入）
self.segmentation_head = nn.Sequential(
nn.Conv2d(hidden_size, 256, kernel_size=1),
nn.Upsample(scale_factor=14, mode='bilinear'), # 518/14=37
nn.Conv2d(256, num_classes, kernel_size=1)
)
# 冻结主干网络（可选解冻部分层以微调）
for param in self.backbone.parameters():
param.requires_grad = False
def forward(self, pixel_values):
# 获取特征 [batch, 37x37+1, hidden_size]
outputs = self.backbone(pixel_values)
last_hidden = outputs.last_hidden_state
# 转换特征形状 [batch, hidden_size, 37, 37]
batch_size = last_hidden.size(0)
features = last_hidden[:, 1:].permute(0, 2, 1) # 移除CLS token
features = features.view(batch_size, -1, 37, 37) # 518/14=37
# 分割头
logits = self.segmentation_head(features)
return logits
# ======================
# 预处理与后处理工具
# ======================
class SegmentationPipeline:
def __init__(self, num_classes=3):
# 确保dinov2_base目录包含正确的preprocessor_config.json
try:
self.feature_extractor = ViTFeatureExtractor.from_pretrained(
"./dinov2_base",
size={"height": 518, "width": 518} # 关键修改
)
except Exception as e:
print(f"Error loading feature extractor: {e}")
raise
self.model = DINOv2ForSegmentation(num_classes)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
self.model.eval()
print(f"Model loaded on {self.device}")
def augment_image(self, image):
"""应用数据增强以提高泛化能力"""
# 亮度
enhancer = ImageEnhance.Brightness(image)
brightness = random.uniform(*AUGMENTATION_PARAMS["brightness_factor"])
image = enhancer.enhance(brightness)
# 对比度
enhancer = ImageEnhance.Contrast(image)
contrast = random.uniform(*AUGMENTATION_PARAMS["contrast_factor"])
image = enhancer.enhance(contrast)
# 色调和饱和度（使用PIL的Color）
enhancer = ImageEnhance.Color(image)
saturation = random.uniform(*AUGMENTATION_PARAMS["saturation_factor"])
image = enhancer.enhance(saturation)
# 旋转（使用PIL的rotate，修正fillmode错误）
angle = random.uniform(*AUGMENTATION_PARAMS["rotation_range"])
image = image.rotate(angle, expand=True, fillcolor=(0, 0, 0)) # 黑色填充
return image
def preprocess(self, image):
if not isinstance(image, Image.Image):
raise ValueError("Input must be a PIL Image")
if image.mode != "RGB":
print(f"Converting image from {image.mode} to RGB")
image = image.convert("RGB")
# 应用数据增强
augmented_image = self.augment_image(image)
# 自动调整到518x518
try:
inputs = self.feature_extractor(
images=augmented_image,
return_tensors="pt",
size={"height": 518, "width": 518}
)
return inputs.pixel_values.to(self.device)
except Exception as e:
print(f"Error in preprocessing: {e}")
raise
def postprocess(self, logits, original_size):
probs = torch.softmax(logits, dim=1)
mask = torch.argmax(probs, dim=1).squeeze().cpu().numpy()
print(f"Mask shape: {mask.shape}, Unique values: {np.unique(mask)}")
print(f"Probability distribution per class: {probs.mean(dim=(0, 2, 3))}") # Debug class probabilities
mask = Image.fromarray(mask.astype(np.uint8))
return mask.resize(original_size, Image.NEAREST)
def predict(self, image, truck_bbox=None):
"""在指定的truck区域内进行预测，如果没有truck区域则使用整个图像"""
if truck_bbox:
# 裁剪图像到truck区域，扩展更大边界以包含边缘
x1, y1, x2, y2 = truck_bbox
padding = 50 # 增加边界以捕捉完整边缘（从 20 增加到 50）
cropped_image = image.crop((
max(0, x1 - padding), max(0, y1 - padding),
min(image.size[0], x2 + padding), min(image.size[1], y2 + padding)
))
else:
cropped_image = image
# 预处理
inputs = self.preprocess(cropped_image)
# 推理
with torch.no_grad():
logits = self.model(inputs)
probs = torch.softmax(logits, dim=1) # Compute probabilities here
print(f"Logits shape: {logits.shape}")
print(f"Logits max: {logits.max()}, min: {logits.min()}")
print(f"Probabilities max: {probs.max()}, min: {probs.min()}")
# 后处理
mask = self.postprocess(logits, cropped_image.size)
# 如果有truck区域，将mask扩展回原图大小
if truck_bbox:
full_mask = Image.new("L", image.size, 0) # 背景为0
adjusted_bbox = (
max(0, x1 - padding), max(0, y1 - padding),
min(image.size[0], x2 + padding), min(image.size[1], y2 + padding)
)
full_mask.paste(mask, adjusted_bbox)
return full_mask
return mask
# ======================
# 使用示例
# ======================
def visualize_and_extract_regions(original_image, mask, save_base_path="output", num_classes=3):
"""
将车斗和覆盖布提取到黑色背景的单独图片上，不再叠加到原图
:param original_image: PIL.Image 原始图片
:param mask: PIL.Image 分割mask
:param save_base_path: 保存路径基础名称
:param num_classes: 分割类别数
"""
# 创建颜色映射 (RGB格式)
colormap = []
# 背景色为黑色（用于单独输出）
colormap.append([0, 0, 0]) # Class 0: Background (black)
# 车斗（红色）、覆盖布（绿色）
colormap.append([255, 0, 0]) # Class 1: Truck bed (red)
colormap.append([0, 255, 0]) # Class 2: Tarp (green)
colormap = np.array(colormap, dtype=np.uint8)
# 将mask转换为数组，确保尺寸匹配
mask_array = np.array(mask.convert("L")) # 确保mask是单通道
original_array = np.array(original_image.convert("RGB"))
height, width = original_array.shape[:2]
# 调整mask尺寸以匹配原图（如果不匹配）
if mask_array.shape != (height, width):
mask_array = np.array(mask.resize((width, height), Image.NEAREST))
# 提取并保存车斗和覆盖布到黑色背景的单独图片，只处理车斗和覆盖布
black_background = np.zeros((height, width, 3), dtype=np.uint8) # 黑色背景
for class_id, class_name in enumerate(CLASS_NAMES[1:], 1): # 跳过背景（0）
class_mask = (mask_array == class_id).astype(np.uint8) * 255 # 二值掩码
# 提取区域到黑色背景
extracted_region = black_background.copy()
for i in range(3): # RGB通道
extracted_region[:, :, i] = black_background[:, :, i] * (1 - class_mask / 255) + \
colormap[class_id, i] * (class_mask / 255)
# 保存提取的区域到黑色背景
extracted_image = Image.fromarray(extracted_region.astype(np.uint8))
extracted_image.save(f"{save_base_path}_{class_name}_on_black.png")
print(f"Extracted {class_name} on black background saved to {save_base_path}_{class_name}_on_black.png")
# 也可以保存原始mask以便检查
mask.save(f"{save_base_path}_mask.png")
print(f"Raw mask saved to {save_base_path}_mask.png")
# 附加：尝试基于颜色后处理以改进结果
color_based_mask = color_based_segmentation(original_image)
if color_based_mask is not None:
visualize_color_based_mask_on_black(original_image, color_based_mask, save_base_path + "_color_based_on_black")
def detect_truck(image_path):
"""
使用YOLOv8检测卡车并返回边界框
"""
try:
# 加载YOLOv8模型
model = YOLO("yolov8n.pt") # 确保yolov8n.pt在当前目录下
# 进行预测
print('yolov8n load successfully~~~')
results = model(image_path, conf=0.3) # confidence threshold 0.3
#print('result:', results)
for result in results:
boxes = result.boxes # 获取检测框
for box in boxes:
cls = int(box.cls[0]) # 类别ID
print('cls:', cls)
if YOLO_CLASSES.get(cls): # == "truck":
print('detect out:', YOLO_CLASSES.get(cls))
x1, y1, x2, y2 = box.xyxy[0].tolist() # 边界框坐标
return (int(x1), int(y1), int(x2), int(y2)) # 返回(x1, y1, x2, y2)
return None # 如果未检测到卡车
except Exception as e:
print(f"Error in YOLOv8 detection: {e}")
return None
def color_based_segmentation(image):
"""
使用颜色阈值分割车斗（多种颜色）和覆盖布（多种颜色），提高泛化能力
"""
try:
# 转换为OpenCV格式 (BGR)
img = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
# 定义更广泛的颜色范围（HSV空间更适合）
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
# 红色范围（车斗，可能为红色、蓝色、黄色等）
lower_red1 = np.array([0, 120, 70]) # 红色范围1
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([170, 120, 70]) # 红色范围2
upper_red2 = np.array([180, 255, 255])
lower_blue = np.array([100, 120, 70]) # 蓝色范围
upper_blue = np.array([130, 255, 255])
lower_yellow = np.array([20, 120, 70]) # 黄色范围
upper_yellow = np.array([40, 255, 255])
mask_red1 = cv2.inRange(hsv, lower_red1, upper_red1)
mask_red2 = cv2.inRange(hsv, lower_red2, upper_red2)
mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)
mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
mask_truck = cv2.bitwise_or(cv2.bitwise_or(mask_red1, mask_red2), cv2.bitwise_or(mask_blue, mask_yellow))
# 绿色范围（覆盖布，可能为绿色、蓝色等）
lower_green = np.array([35, 40, 40])
upper_green = np.array([85, 255, 255])
lower_blue_tarp = np.array([100, 40, 40]) # 蓝色覆盖布
upper_blue_tarp = np.array([130, 255, 255])
mask_green = cv2.inRange(hsv, lower_green, upper_green)
mask_blue_tarp = cv2.inRange(hsv, lower_blue_tarp, upper_blue_tarp)
mask_tarp = cv2.bitwise_or(mask_green, mask_blue_tarp)
# 合并掩码：0=背景, 1=车斗, 2=覆盖布
combined_mask = np.zeros((hsv.shape[0], hsv.shape[1]), dtype=np.uint8)
combined_mask[mask_truck > 0] = 1 # 车斗
combined_mask[mask_tarp > 0] = 2 # 覆盖布
# 应用形态学操作去除噪声并填充小孔
kernel = np.ones((5, 5), np.uint8)
combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel)
combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)
return Image.fromarray(combined_mask)
except Exception as e:
print(f"Error in color-based segmentation: {e}")
return None
def visualize_color_based_mask_on_black(original_image, mask, save_path):
"""
可视化基于颜色的分割结果到黑色背景
"""
colormap = np.array([
[0, 0, 0], # 背景 (黑)
[255, 0, 0], # 车斗 (红)
[0, 255, 0] # 覆盖布 (绿)
], dtype=np.uint8)
mask_array = np.array(mask)
height, width = mask_array.shape
black_background = np.zeros((height, width, 3), dtype=np.uint8)
# 提取区域到黑色背景
for class_id in [1, 2]: # 只处理车斗和覆盖布
class_mask = (mask_array == class_id).astype(np.uint8) * 255
for i in range(3): # RGB通道
black_background[:, :, i] = black_background[:, :, i] * (1 - class_mask / 255) + \
colormap[class_id, i] * (class_mask / 255)
extracted_image = Image.fromarray(black_background.astype(np.uint8))
extracted_image.save(save_path + ".png")
print(f"Color-based visualization on black background saved to {save_path}.png")
def enhance_segmentation(mask_array, truck_region, num_classes):
"""
增强分割结果以捕捉车斗的完整轮廓，处理边缘和颜色遮挡
"""
# 应用更强的形态学操作以填补空洞和捕捉边缘
kernel_large = np.ones((20, 20), np.uint8) # 增大内核以捕捉完整边缘（从 15 增加到 20）
kernel_small = np.ones((3, 3), np.uint8) # 用于细化边缘
# 膨胀以捕捉完整轮廓
dilated_mask = cv2.dilate(mask_array, kernel_large, iterations=4) # 增加迭代次数以捕捉更多边缘
# 腐蚀以去除噪声，保持边界
eroded_mask = cv2.erode(dilated_mask, kernel_large, iterations=1)
# 闭运算填补空洞
closed_mask = cv2.morphologyEx(eroded_mask, cv2.MORPH_CLOSE, kernel_large)
# 开运算去除小噪声
opened_mask = cv2.morphologyEx(closed_mask, cv2.MORPH_OPEN, kernel_small)
# 细化边缘：使用更敏感的 Canny 边缘检测
edges = cv2.Canny((opened_mask == 1).astype(np.uint8) * 255, 30, 100) # 降低阈值以捕捉更多细小边缘
dilated_edges = cv2.dilate(edges, kernel_small, iterations=3) # 增加膨胀以连接边缘
# 使用颜色和形状信息进一步优化车斗区域
hsv_truck = cv2.cvtColor(truck_region, cv2.COLOR_RGB2HSV)
height, width = opened_mask.shape
# 查找车斗和覆盖布的轮廓
truck_bed_mask = (opened_mask == 1).astype(np.uint8) * 255
tarp_mask = (opened_mask == 2).astype(np.uint8) * 255
# 轮廓检测，填充完整轮廓
contours_truck, _ = cv2.findContours(truck_bed_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contours_tarp, _ = cv2.findContours(tarp_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# 创建增强后的掩码
enhanced_mask = np.zeros_like(opened_mask, dtype=np.uint8)
# 填充车斗轮廓，确保完整性并结合边缘
if contours_truck:
# 找到最大的轮廓（假设车斗是最大的区域）
largest_contour = max(contours_truck, key=cv2.contourArea)
cv2.drawContours(enhanced_mask, [largest_contour], -1, 1, thickness=cv2.FILLED)
# 结合边缘信息，填补细小边界
enhanced_mask[dilated_edges > 0] = 1 # 将检测到的边缘区域标记为车斗
else:
# 如果没有检测到轮廓，使用膨胀后的区域
enhanced_mask[opened_mask == 1] = 1
# 填充覆盖布轮廓
for contour in contours_tarp:
cv2.drawContours(enhanced_mask, [contour], -1, 2, thickness=cv2.FILLED)
# 确保掩码值在有效范围内
enhanced_mask = np.clip(enhanced_mask, 0, num_classes - 1)
# 额外处理：如果车斗区域有小断裂或边缘缺失，使用连通性分析填补
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(truck_bed_mask, connectivity=8)
if num_labels > 1: # 如果有多个连通区域
largest_area = 0
largest_label = 0
for label in range(1, num_labels): # 跳过背景（标签0）
area = stats[label, cv2.CC_STAT_AREA]
if area > largest_area:
largest_area = area
largest_label = label
enhanced_mask[labels == largest_label] = 1 # 保留最大连通区域作为车斗
# 再次细化边缘，确保完整性
final_edges = cv2.Canny((enhanced_mask == 1).astype(np.uint8) * 255, 20, 80) # 进一步降低阈值捕捉边缘
enhanced_mask[final_edges > 0] = 1 # 填补边缘
# 边界扩展：额外膨胀以确保边缘完整
final_dilated = cv2.dilate((enhanced_mask == 1).astype(np.uint8) * 255, kernel_small, iterations=2)
enhanced_mask[final_dilated > 0] = 1 # 扩展车斗边缘
return enhanced_mask
if __name__ == "__main__":
# 初始化 pipeline
num_classes = 3 # 背景, 车斗, 覆盖布
try:
pipeline = SegmentationPipeline(num_classes=num_classes)
# 加载原始图像
image_path = "trunk.jpg"
original_image = Image.open(image_path)
# 使用YOLOv8检测卡车
truck_bbox = detect_truck(image_path)
if truck_bbox:
print(f"Detected truck bounding box: {truck_bbox}")
# 进行分割
segmentation_mask = pipeline.predict(original_image, truck_bbox)
# 提取并可视化车斗和覆盖布到黑色背景
visualize_and_extract_regions(
original_image=original_image,
mask=segmentation_mask,
save_base_path="segmentation_output",
num_classes=num_classes
)
else:
print("No truck detected, using full image for segmentation")
except Exception as e:
print(f"Error in main execution: {e}")- 扩展 YOLO 边界框（padding）：
-
在 predict 函数中，将 padding 从 20 增加到 50，确保 YOLO 检测的边界框包含车斗的完整边缘，减少因边界框过紧而丢失边缘的可能性。
-
- 增强边缘检测（Canny 和形态学操作）：
-
增加额外膨胀步骤（final_dilated 使用 3x3 内核，迭代 2 次）以扩展车斗边缘，确保边缘完整。
-
优化 Canny 边缘检测，降低阈值（从 30, 100 调整到 20, 80 再到 20, 80），确保捕捉更多细小边缘。
-
增大形态学操作的内核大小（kernel_large 从 15x15 增加到 20x20），并增加膨胀迭代次数（从 3 增加到 4），更好地捕捉车斗边缘
-
- 结合颜色和形状信息：
-
保留颜色分割（color_based_segmentation）和轮廓检测（findContours），结合最大连通区域和边缘信息，确保车斗的完整轮廓被分割，即使有光影变化或遮挡。
-
使用连通性分析（connectedComponentsWithStats）保留最大连通区域，填补小断裂或边缘缺失。
-
- 预期效果：
-
车斗的完整轮廓（包括边缘）将被标记为红色，输出到 segmentation_output_truck_bed_on_black.png，即使边缘细小或有光影变化。
-
覆盖布（绿色）仍会被正确分割到 segmentation_output_tarp_on_black.png，但不会干扰车斗的完整性。
-
调试和下一步
- 运行优化后的代码，检查输出 segmentation_output_truck_bed_on_black.png 和 segmentation_output_enhanced_mask.png，确保车斗边缘被完整分割。
- 如果车斗边缘仍不完整，尝试：
- 进一步增加 padding（如 70 或更高）。
- 调整 enhance_segmentation 中的 kernel_large（如 25x25）、迭代次数，或 Canny 阈值（试试 10, 60）。
在语义分割的基础上，下面是判断车斗上方空间的范围何雨覆的位置
分析当前分割结果
- 车斗（红色）：从图片看，车斗的轮廓已基本完整，但可能有细小噪声或不规则边缘。
- 覆盖布（绿色）：覆盖布分布在车斗顶部，但可能有间断或未完全覆盖的部分。
- 目标：
- 计算车斗的顶部投影面积（假设为车斗的完整轮廓面积）。
- 计算覆盖布的面积（绿色区域面积）。
- 比较覆盖布面积与车斗面积，判断覆盖率是否达到 80%。
这里给出算法代码
import cv2
import numpy as np
from PIL import Image
# 加载图像
image_path = "truck_seg.png"
image = cv2.imread(image_path)
# 转换为 RGB 格式（OpenCV 默认 BGR）
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 转换为 HSV 格式以便颜色分割
hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
# 定义颜色范围（基于红色车斗和绿色雨覆）
# 红色范围（车斗，可能为红色、蓝色、黄色等）
lower_red1 = np.array([0, 120, 70]) # 红色范围1
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([170, 120, 70]) # 红色范围2
upper_red2 = np.array([180, 255, 255])
lower_blue = np.array([100, 120, 70]) # 蓝色范围
upper_blue = np.array([130, 255, 255])
lower_yellow = np.array([20, 120, 70]) # 黄色范围
upper_yellow = np.array([40, 255, 255])
# 绿色范围（雨覆，可能为绿色、蓝色等）
lower_green = np.array([35, 40, 40])
upper_green = np.array([85, 255, 255])
lower_blue_tarp = np.array([100, 40, 40]) # 蓝色雨覆
upper_blue_tarp = np.array([130, 255, 255])
# 颜色分割
mask_red1 = cv2.inRange(hsv, lower_red1, upper_red1)
mask_red2 = cv2.inRange(hsv, lower_red2, upper_red2)
mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)
mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
mask_truck = cv2.bitwise_or(cv2.bitwise_or(mask_red1, mask_red2), cv2.bitwise_or(mask_blue, mask_yellow))
mask_green = cv2.inRange(hsv, lower_green, upper_green)
mask_blue_tarp = cv2.inRange(hsv, lower_blue_tarp, upper_blue_tarp)
mask_tarp = cv2.bitwise_or(mask_green, mask_blue_tarp)
# 合并掩码：0=背景, 1=车斗, 2=雨覆
combined_mask = np.zeros((hsv.shape[0], hsv.shape[1]), dtype=np.uint8)
combined_mask[mask_truck > 0] = 1 # 车斗
combined_mask[mask_tarp > 0] = 2 # 雨覆
# 应用形态学操作去除噪声并填充空洞
kernel = np.ones((5, 5), np.uint8)
closed_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel)
opened_mask = cv2.morphologyEx(closed_mask, cv2.MORPH_OPEN, kernel)
# 增强车斗轮廓，捕捉完整边缘
truck_bed_mask = (opened_mask == 1).astype(np.uint8) * 255
kernel_large = np.ones((20, 20), np.uint8)
dilated_truck = cv2.dilate(truck_bed_mask, kernel_large, iterations=4)
eroded_truck = cv2.erode(dilated_truck, kernel_large, iterations=1)
closed_truck = cv2.morphologyEx(eroded_truck, cv2.MORPH_CLOSE, kernel_large)
# 增强雨覆轮廓
tarp_mask = (opened_mask == 2).astype(np.uint8) * 255
dilated_tarp = cv2.dilate(tarp_mask, kernel_large, iterations=2)
closed_tarp = cv2.morphologyEx(dilated_tarp, cv2.MORPH_CLOSE, kernel_large)
# 更新增强后的掩码
enhanced_mask = np.zeros_like(opened_mask, dtype=np.uint8)
enhanced_mask[closed_truck > 0] = 1 # 车斗
enhanced_mask[closed_tarp > 0] = 2 # 雨覆
# 计算面积（像素数）
truck_bed_area = np.sum(enhanced_mask == 1) # 车斗面积
tarp_area = np.sum(enhanced_mask == 2) # 雨覆面积
# 计算覆盖率
if truck_bed_area > 0:
coverage_ratio = (tarp_area / truck_bed_area) * 100 # 覆盖率（百分比）
else:
coverage_ratio = 0 # 如果车斗面积为0，覆盖率设为0
# 打印面积和覆盖率
print(f"Truck bed area (pixels): {truck_bed_area}")
print(f"Tarp area (pixels): {tarp_area}")
print(f"Tarp coverage ratio: {coverage_ratio:.2f}%")
print(f"Coverage meets 80% requirement: {'Yes' if coverage_ratio >= 80 else 'No'}")
# 可视化：提取车斗和雨覆到黑色背景
height, width = enhanced_mask.shape
black_background = np.zeros((height, width, 3), dtype=np.uint8)
# 提取车斗（红色）
truck_mask = (enhanced_mask == 1).astype(np.uint8) * 255
black_background[truck_mask > 0] = [255, 0, 0] # 红色
# 提取雨覆（绿色）
tarp_mask = (enhanced_mask == 2).astype(np.uint8) * 255
black_background[tarp_mask > 0] = [0, 255, 0] # 绿色
# 保存结果
cv2.imwrite("truck_bed_on_black.png", cv2.cvtColor(black_background, cv2.COLOR_RGB2BGR))
print(f"Extracted truck bed on black background saved to truck_bed_on_black.png")
cv2.imwrite("tarp_on_black.png", cv2.cvtColor(black_background, cv2.COLOR_RGB2BGR))
print(f"Extracted tarp on black background saved to tarp_on_black.png")
# 保存增强后的掩码以便检查
cv2.imwrite("enhanced_mask.png", enhanced_mask)
print(f"Enhanced mask saved to enhanced_mask.png")- 图像加载和颜色分割：
-
加载 truck_seg.png，转换为 HSV 颜色空间。
-
使用预定义的红色（车斗）和绿色（雨覆）范围进行颜色分割，规则与之前一致，支持多种颜色变体（红、蓝、黄；绿、蓝）。
-
- 形态学操作和轮廓增强：
- 应用闭运算（MORPH_CLOSE）填补空洞，开运算（MORPH_OPEN）去除噪声。
- 增强车斗和雨覆轮廓，使用更大的内核（20x20）和更多迭代次数（车斗 4 次，雨覆 2 次）以捕捉完整边缘。
- 面积计算：
- 使用 np.sum 统计车斗（值为 1）和雨覆（值为 2）的像素数，单位为像素。
- 计算覆盖率：tarp_area / truck_bed_area * 100，判断是否 ≥ 80%。
- 可视化输出：
- 将车斗（红色）和雨覆（绿色）提取到黑色背景，分别保存为 truck_bed_on_black.png 和 tarp_on_black.png。
- 保存增强后的掩码 enhanced_mask.png 以便检查。
经过检测覆盖率没有达到80%
上面图像分割还需要再优化下，毕竟存在把非雨覆物体错误识别的地方，但整体思路可以分享出来供大家借鉴
