DINOv2 + yolov8 + opencv 检测卡车的可拉拽雨覆是否完全覆盖

devtools/2025/2/28 3:15:09/

最近是接了一个需求咨询图像处理类的，甲方要在卡车过磅的地方装一个摄像头用检测卡车的车斗雨覆是否完全， 让我大致理了下需求并对技术核心做下预研究

开发一套图像处理软件，能够实时监控经过的卡车并判断其车斗的雨覆状态。

系统需具备以下功能：

  1. 图像采集：通过高分辨率摄像头采集卡车经过时的图像。
  2. 图像处理：对采集的图像进行处理，识别车斗及雨覆的具体位置。
  3. 状态判断：
    • 判断雨覆是否完全覆盖车斗。
    • 在雨覆未完全覆盖时，生成警报或提示信息。

例如这样的图片，检测上方雨覆是否完全遮盖住

需求分析

图像采集是视频流接收，确保摄像头支持所选协议，并具备高清分辨率（至少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",
}

  • 作用：
    • CLASS_NAMES：定义语义分割的类别（背景、车斗、覆盖布），用于 DINOv2 模型的输出和后续处理。
    • AUGMENTATION_PARAMS：设置数据增强参数，模拟亮度、对比度、旋转、色调和饱和度的变化，提高模型对不同光照、角度和颜色的泛化能力。
    • YOLO_CLASSES：定义 YOLOv8 模型的类别映射，指定类别 ID（如 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 = Falsedef forward(self, pixel_values):# 从 DINOv2 提取特征并通过分割头生成分割结果outputs = self.backbone(pixel_values)last_hidden = outputs.last_hidden_statefeatures = 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 imagedef preprocess(self, image):# 预处理图像：转换为 RGB，应用增强，调整到 518x518if 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 maskdef 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 = imageinputs = 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_maskreturn 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) * 255extracted_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] = 1combined_mask[mask_tarp > 0] = 2kernel = 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 的补充。
      颜色分割可视化（visualize_color_based_mask_on_black）
    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) * 255for 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__":# 初始化 pipelinenum_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}")raisehidden_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=37nn.Conv2d(256, num_classes, kernel_size=1))# 冻结主干网络（可选解冻部分层以微调）for param in self.backbone.parameters():param.requires_grad = Falsedef 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 tokenfeatures = 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.jsontry:self.feature_extractor = ViTFeatureExtractor.from_pretrained("./dinov2_base",size={"height": 518, "width": 518}  # 关键修改)except Exception as e:print(f"Error loading feature extractor: {e}")raiseself.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 imagedef 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)# 自动调整到518x518try: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}")raisedef 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 probabilitiesmask = 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_bboxpadding = 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 hereprint(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)  # 背景为0adjusted_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_maskreturn 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])  # 类别IDprint('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 Nonedef 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])  # 红色范围1upper_red1 = np.array([10, 255, 255])lower_red2 = np.array([170, 120, 70])  # 红色范围2upper_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 Nonedef 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.shapeblack_background = np.zeros((height, width, 3), dtype=np.uint8)# 提取区域到黑色背景for class_id in [1, 2]:  # 只处理车斗和覆盖布class_mask = (mask_array == class_id).astype(np.uint8) * 255for 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) * 255tarp_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 = 0largest_label = 0for label in range(1, num_labels):  # 跳过背景（标签0）area = stats[label, cv2.CC_STAT_AREA]if area > largest_area:largest_area = arealargest_label = labelenhanced_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_maskif __name__ == "__main__":# 初始化 pipelinenum_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）。

    在语义分割的基础上，下面是判断车斗上方空间的范围何雨覆的位置

    分析当前分割结果

    • 车斗（红色）：从图片看，车斗的轮廓已基本完整，但可能有细小噪声或不规则边缘。
    • 覆盖布（绿色）：覆盖布分布在车斗顶部，但可能有间断或未完全覆盖的部分。
    • 目标：
      1. 计算车斗的顶部投影面积（假设为车斗的完整轮廓面积）。
      2. 计算覆盖布的面积（绿色区域面积）。
      3. 比较覆盖布面积与车斗面积，判断覆盖率是否达到 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%

    上面图像分割还需要再优化下，毕竟存在把非雨覆物体错误识别的地方，但整体思路可以分享出来供大家借鉴


    http://www.ppmy.cn/devtools/163233.html

    相关文章

    C++关键字之mutable

    C++关键字之mutable

    1.介绍 在C中，mutable是一个关键字，用于修饰类的成员变量。它的主要作用是允许在常量成员函数或常量对象中修改被标记为mutable的成员变量。通常情况下，常量成员函数不能修改类的成员变量，但有些情况下，某些成员变量的…
    阅读更多...
    Redis分布式缓存面试题

    Redis分布式缓存面试题

    为什么使用分布式缓存？ 1. 提升性能 降低延迟：将数据缓存在离应用更近的地方，减少数据访问时间。减轻数据库压力：缓存频繁访问的数据，减少对后端数据库的请求，提升系统响应速度。 2. 扩展性 水平扩展&a…
    阅读更多...
    Linux基础开发工具的使用（apt、vim、gcc、g++、gdb、make、makefile）

    Linux基础开发工具的使用（apt、vim、gcc、g++、gdb、make、makefile）

    Linux软件包管理器–apt Linux安装软件的方式 在Linux下安装软件的方法有以下三种： 下载到程序的源代码，自己编译出可执行程序获取deb安装包、然后使用dpkg命令安装。（不解决依赖关系）通过apt进行安装软件。 小知识点&#xf…
    阅读更多...
    云计算如何解决延迟问题？

    云计算如何解决延迟问题？

    在云计算中，延迟（latency）指的是从请求发出到收到响应之间的时间间隔。延迟过高可能会严重影响用户体验，特别是在需要实时响应的应用中，如在线游戏、视频流、金融交易等。云计算服务如何解决延迟问题，通常依…
    阅读更多...
    汽车结构胶仿真模型MAT_169材料卡片的制作

    汽车结构胶仿真模型MAT_169材料卡片的制作

    随着汽车轻量化技术的发展，车身所用材料呈现出多样化的趋势，由于异种材料之间的物理、化学和力学性能方面存在较大差异，因此多材料轻量化车身对连接技术提出了新的挑战。 传统点焊连接由于技术瓶颈和成本的原因，无法广泛应用于异种…
    阅读更多...
    力扣hot100 —— 电话号码字母组合； 子集 （非回溯做法）简单易懂

    力扣hot100 —— 电话号码字母组合； 子集 （非回溯做法）简单易懂

    由于博主对回溯也不是很熟悉，这里提出一种简单易懂的解法（有点暴力） 解题思路： 每个数字对应有自己的字母串； 首先遍历将每个字母存入也就是 res{{a},{b},{c}} 然后遍历后续数子对应的字母，让每个字母与…
    阅读更多...
    【自学笔记】Vue基础知识点总览-持续更新

    【自学笔记】Vue基础知识点总览-持续更新

    提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档 文章目录 引言Vue基础知识概览1. Vue实例2. 模板语法3. 计算属性4. 事件处理 总结 引言 Vue.js是一个构建用户界面的渐进式框架，以其简洁的API和易于上手的特点受…
    阅读更多...
    【鸿蒙应用开发】性能优化

    【鸿蒙应用开发】性能优化

    渲染方面 Repeat：可复用的循环渲染 Repeat 一般会用于取代 ForEach，相较后者具有更强的渲染性能，Repeat 具有两种工作模式： non-virtualScroll 模式 在初始化页面时就加载列表中的全部子组件。 相比于ForEach，具有…
    阅读更多...
    最新文章

    Copyright @ 2022~2023