2018年淘宝客网站怎么做,电子商务网站设计岗位主要是,提供企业网站建设方案,科学小制作最近是接了一个需求咨询图像处理类的#xff0c;甲方要在卡车过磅的地方装一个摄像头用检测卡车的车斗雨覆是否完全#xff0c; 让我大致理了下需求并对技术核心做下预研究
开发一套图像处理软件#xff0c;能够实时监控经过的卡车并判断其车斗的雨覆状态。
系统需具备以下…最近是接了一个需求咨询图像处理类的甲方要在卡车过磅的地方装一个摄像头用检测卡车的车斗雨覆是否完全 让我大致理了下需求并对技术核心做下预研究
开发一套图像处理软件能够实时监控经过的卡车并判断其车斗的雨覆状态。
系统需具备以下功能
图像采集通过高分辨率摄像头采集卡车经过时的图像。图像处理对采集的图像进行处理识别车斗及雨覆的具体位置。状态判断 判断雨覆是否完全覆盖车斗。在雨覆未完全覆盖时生成警报或提示信息。
例如这样的图片检测上方雨覆是否完全遮盖住 需求分析
图像采集是视频流接收确保摄像头支持所选协议并具备高清分辨率至少1080p以提高图像识别的准确性。摄像头应具备良好的低光性能以适应不同的环境光照条件。
使用开源媒体框架如 FFmpeg 或 GStreamer来接收和处理视频流。实现视频流的解码提取每一帧图像供后续处理。
所以在技术预研上直接从图片开始 目测了下需要使用图像语义分割再在分割图像基础上再计算雨覆的遮盖率会使用到的工具具体有
导入必要的库用于深度学习PyTorch、图像处理PIL、OpenCV、可视化Matplotlib、以及 YOLOv8Ultralytics目标检测和 DINOv2Transformers语义分割。
PyTorch的作用 提供核心深度学习功能如 torch.nn 用于定义 DINOv2 分割模型DINOv2ForSegmentation 类。通过 torch.device 确定设备CPU 或 CUDA将模型和数据加载到 GPUself.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 绘制图表仅用于文件保存 UltralyticsYOLOv8的作用 加载预训练的 yolov8n.pt 模型YOLO(yolov8n.pt)检测图像中的卡车detect_truck 函数。返回卡车的边界框box.xyxy用于裁剪图像区域供 DINOv2 语义分割确保仅在 truck 区域内分割车斗和覆盖布。设置置信度阈值conf0.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_classes3, 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_size1),nn.Upsample(scale_factor14, modebilinear),nn.Conv2d(256, num_classes, kernel_size1))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 的语义分割模型使用预训练的 ViTVision Transformer作为主干网络冻结其参数以减少计算量。添加自定义分割头segmentation_head将 37x37 的特征图上采样并生成 3 类的分割结果背景、车斗、覆盖布。forward 方法处理输入图像518x518输出分割 logits。
分割管道SegmentationPipeline
class SegmentationPipeline:def __init__(self, num_classes3):# 初始化 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, expandTrue, 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(imagesaugmented_image, return_tensorspt, size{height: 518, width: 518})return inputs.pixel_values.to(self.device)def postprocess(self, logits, original_size):# 后处理softmax 转换为概率取最大值生成掩码调整回原图大小probs torch.softmax(logits, dim1)mask torch.argmax(probs, dim1).squeeze().cpu().numpy()mask Image.fromarray(mask.astype(np.uint8)).resize(original_size, Image.NEAREST)return maskdef predict(self, image, truck_bboxNone):# 在 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, dim1)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 边界框裁剪图像进行分割支持扩展边界padding20以捕捉边缘。
可视化和提取区域visualize_and_extract_regions
def visualize_and_extract_regions(original_image, mask, save_base_pathoutput, num_classes3):# 创建黑色背景提取车斗和覆盖布保存到单独图片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), dtypenp.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, conf0.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检测图像中的卡车类别 truckID 7。设置置信度阈值conf0.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], dtypenp.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), dtypenp.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_classesnum_classes)# 加载原始图像image_path trunk.jpgoriginal_image Image.open(image_path)# 使用YOLOv8检测卡车truck_bbox detect_truck(image_path)if truck_bbox:print(fDetected truck bounding box: {truck_bbox})# 进行分割segmentation_mask pipeline.predict(original_image, truck_bbox)# 提取并可视化车斗和覆盖布到黑色背景visualize_and_extract_regions(original_imageoriginal_image,masksegmentation_mask,save_base_pathsegmentation_output,num_classesnum_classes)else:print(No truck detected, using full image for segmentation)except Exception as e:print(fError 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_classes3, model_name./dinov2_base/):super().__init__()# 加载 DINOv2 主干网络try:self.backbone ViTModel.from_pretrained(model_name)except Exception as e:print(fError loading DINOv2 model: {e})raisehidden_size self.backbone.config.hidden_size# 分割头调整适配518输入self.segmentation_head nn.Sequential(nn.Conv2d(hidden_size, 256, kernel_size1),nn.Upsample(scale_factor14, modebilinear), # 518/1437nn.Conv2d(256, num_classes, kernel_size1))# 冻结主干网络可选解冻部分层以微调for param in self.backbone.parameters():param.requires_grad Falsedef forward(self, pixel_values):# 获取特征 [batch, 37x371, 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/1437# 分割头logits self.segmentation_head(features)return logits#
# 预处理与后处理工具
#
class SegmentationPipeline:def __init__(self, num_classes3):# 确保dinov2_base目录包含正确的preprocessor_config.jsontry:self.feature_extractor ViTFeatureExtractor.from_pretrained(./dinov2_base,size{height: 518, width: 518} # 关键修改)except Exception as e:print(fError 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(fModel 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的Colorenhancer 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, expandTrue, 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(fConverting image from {image.mode} to RGB)image image.convert(RGB)# 应用数据增强augmented_image self.augment_image(image)# 自动调整到518x518try:inputs self.feature_extractor(imagesaugmented_image,return_tensorspt,size{height: 518, width: 518})return inputs.pixel_values.to(self.device)except Exception as e:print(fError in preprocessing: {e})raisedef postprocess(self, logits, original_size):probs torch.softmax(logits, dim1)mask torch.argmax(probs, dim1).squeeze().cpu().numpy()print(fMask shape: {mask.shape}, Unique values: {np.unique(mask)})print(fProbability 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_bboxNone):在指定的truck区域内进行预测如果没有truck区域则使用整个图像if truck_bbox:# 裁剪图像到truck区域扩展更大边界以包含边缘x1, y1, x2, y2 truck_bboxpadding 50 # 增加边界以捕捉完整边缘从 20 增加到 50cropped_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, dim1) # Compute probabilities hereprint(fLogits shape: {logits.shape})print(fLogits max: {logits.max()}, min: {logits.min()})print(fProbabilities 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_pathoutput, num_classes3):将车斗和覆盖布提取到黑色背景的单独图片上不再叠加到原图: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, dtypenp.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), dtypenp.uint8) # 黑色背景for class_id, class_name in enumerate(CLASS_NAMES[1:], 1): # 跳过背景0class_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(fExtracted {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(fRaw 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, conf0.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(fError 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]), dtypenp.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(fError 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] # 覆盖布 (绿)], dtypenp.uint8)mask_array np.array(mask)height, width mask_array.shapeblack_background np.zeros((height, width, 3), dtypenp.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(fColor-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 增加到 20kernel_small np.ones((3, 3), np.uint8) # 用于细化边缘# 膨胀以捕捉完整轮廓dilated_mask cv2.dilate(mask_array, kernel_large, iterations4) # 增加迭代次数以捕捉更多边缘# 腐蚀以去除噪声保持边界eroded_mask cv2.erode(dilated_mask, kernel_large, iterations1)# 闭运算填补空洞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, iterations3) # 增加膨胀以连接边缘# 使用颜色和形状信息进一步优化车斗区域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, dtypenp.uint8)# 填充车斗轮廓确保完整性并结合边缘if contours_truck:# 找到最大的轮廓假设车斗是最大的区域largest_contour max(contours_truck, keycv2.contourArea)cv2.drawContours(enhanced_mask, [largest_contour], -1, 1, thicknesscv2.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, thicknesscv2.FILLED)# 确保掩码值在有效范围内enhanced_mask np.clip(enhanced_mask, 0, num_classes - 1)# 额外处理如果车斗区域有小断裂或边缘缺失使用连通性分析填补num_labels, labels, stats, _ cv2.connectedComponentsWithStats(truck_bed_mask, connectivity8)if num_labels 1: # 如果有多个连通区域largest_area 0largest_label 0for label in range(1, num_labels): # 跳过背景标签0area 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, iterations2)enhanced_mask[final_dilated 0] 1 # 扩展车斗边缘return enhanced_maskif __name__ __main__:# 初始化 pipelinenum_classes 3 # 背景, 车斗, 覆盖布try:pipeline SegmentationPipeline(num_classesnum_classes)# 加载原始图像image_path trunk.jpgoriginal_image Image.open(image_path)# 使用YOLOv8检测卡车truck_bbox detect_truck(image_path)if truck_bbox:print(fDetected truck bounding box: {truck_bbox})# 进行分割segmentation_mask pipeline.predict(original_image, truck_bbox)# 提取并可视化车斗和覆盖布到黑色背景visualize_and_extract_regions(original_imageoriginal_image,masksegmentation_mask,save_base_pathsegmentation_output,num_classesnum_classes)else:print(No truck detected, using full image for segmentation)except Exception as e:print(fError 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]), dtypenp.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, iterations4)
eroded_truck cv2.erode(dilated_truck, kernel_large, iterations1)
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, iterations2)
closed_tarp cv2.morphologyEx(dilated_tarp, cv2.MORPH_CLOSE, kernel_large)# 更新增强后的掩码
enhanced_mask np.zeros_like(opened_mask, dtypenp.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(fTruck bed area (pixels): {truck_bed_area})
print(fTarp area (pixels): {tarp_area})
print(fTarp coverage ratio: {coverage_ratio:.2f}%)
print(fCoverage meets 80% requirement: {Yes if coverage_ratio 80 else No})# 可视化提取车斗和雨覆到黑色背景
height, width enhanced_mask.shape
black_background np.zeros((height, width, 3), dtypenp.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(fExtracted 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(fExtracted tarp on black background saved to tarp_on_black.png)# 保存增强后的掩码以便检查
cv2.imwrite(enhanced_mask.png, enhanced_mask)
print(fEnhanced 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%
上面图像分割还需要再优化下毕竟存在把非雨覆物体错误识别的地方但整体思路可以分享出来供大家借鉴
