最近是接了一個需求咨詢圖像處理類的，甲方要在卡車過磅的地方裝一個攝像頭用檢測卡車的車斗雨覆是否完全， 讓我大致理了下需求并對技術核心做下預研究

開發一套圖像處理軟件，能夠實時監控經過的卡車并判斷其車斗的雨覆狀態。

系統需具備以下功能：

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%

上面圖像分割還需要再優化下，畢竟存在把非雨覆物體錯誤識別的地方，但整體思路可以分享出來供大家借鑒