Python案例如何实现场景分割？

wen python案例 2026-06-16 230

本文目录导读：

Python案例如何实现场景分割？

方法一：使用 SegFormer（最推荐，高精度）
方法二：使用 SAM（Segment Anything Model，万物分割）
方法三：使用 DeepLabv3+（轻量级、速度快）
方法四：基于 Contour 的传统图像分割（无深度学习）
总结与选择建议

实现Python场景分割（图像分割）通常使用深度学习模型，最主流的方法是基于SegFormer（如Mask2Former）、DeepLabv3+ 或 SAM（Segment Anything Model）。

以下我将提供4种常见且实用的实现方式，从最强大的到最轻量的，并附上完整代码案例。

前置准备：安装依赖

pip install torch torchvision
pip install transformers     # 用于SegFormer
pip install opencv-python    # 用于图像处理
pip install matplotlib       # 用于显示结果

使用 SegFormer（最推荐，高精度）

SegFormer是NVIDIA提出的高效Transformer分割模型,适合语义分割场景。

代码实现

import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
def segment_with_segformer(image_path,save_path='seg_result.jpg'):
    # 1. 加载预训练模型和处理器
    # 选择模型：b0最快但精度一般，b5最慢但精度极高
    model_name = "nvidia/segformer-b0-finetuned-ade-512-512"
    processor = SegformerImageProcessor.from_pretrained(model_name)
    model = SegformerForSemanticSegmentation.from_pretrained(model_name)
    # 2. 读取图像
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 3. 预处理
    inputs = processor(images=image_rgb, return_tensors="pt")
    # 4. 推理
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits  # shape: [1, 150, 128, 128]
        # 上采样到原图尺寸
        upsampled_logits = torch.nn.functional.interpolate(
            logits,
            size=image_rgb.shape[:2],
            mode="bilinear",
            align_corners=False
        )
        # 获取每个像素的类别索引
        predicted_map = upsampled_logits.argmax(dim=1).squeeze(0).cpu().numpy()
    # 5. 可视化（将预测的类别映射为颜色）
    # ADE20K数据集的默认颜色映射（150类）
    color_map = np.random.randint(0, 255, (150, 3), dtype=np.uint8)
    segmented_image = color_map[predicted_map]
    # 6. 显示结果
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.imshow(image_rgb)
    plt.title("原始图像")
    plt.axis('off')
    plt.subplot(1, 2, 2)
    plt.imshow(segmented_image)
    plt.title("语义分割结果")
    plt.axis('off')
    plt.savefig(save_path)
    plt.show()
    return predicted_map, segmented_image
# 使用示例
image_path = "/path/to/your/image.jpg"  # 替换为你的图片路径
segment_with_segformer(image_path)

注：ADE20K数据集包含150个语义类别（道路、树木、天空、人等）。

使用 SAM（Segment Anything Model，万物分割）

如果你需要分割出图像中的每个物体（无需预先定义类别），SAM是最强的选择。

import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
from segment_anything import sam_model_registry, SamPredictor
def segment_with_sam(image_path, save_path='sam_result.jpg'):
    # 1. 加载SAM模型（需要先下载权重文件）
    # 下载地址：https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
    sam_checkpoint = "/path/to/sam_vit_h_4b8939.pth"  # 替换为实际路径
    device = "cuda" if torch.cuda.is_available() else "cpu"
    sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
    sam.to(device=device)
    predictor = SamPredictor(sam)
    # 2. 读取图像
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 3. 设置图像
    predictor.set_image(image_rgb)
    # 4. 自动生成全图掩码（不需要输入点）
    # 使用`generate`模式或使用密集网格点
    h, w = image_rgb.shape[:2]
    point_grid_size = 32  # 网格密度
    points_per_side = point_grid_size
    from segment_anything import SamAutomaticMaskGenerator
    mask_generator = SamAutomaticMaskGenerator(sam)
    masks = mask_generator.generate(image_rgb)
    # 5. 可视化所有分割结果
    plt.figure(figsize=(12, 12))
    plt.imshow(image_rgb)
    for i, mask_data in enumerate(masks):
        mask = mask_data['segmentation']
        # 随机颜色
        color = np.random.rand(3)
        overlay = np.zeros_like(image_rgb, dtype=np.float32)
        overlay[mask] = color * 255
        plt.imshow(overlay, alpha=0.4)
    plt.axis('off')
    plt.savefig(save_path)
    plt.show()
    return masks
# 使用示例
image_path = "/path/to/your/image.jpg"
segment_with_sam(image_path)

使用 DeepLabv3+（轻量级、速度快）

适合实时应用或移动端部署,ResNet101作为backbone。

import torch
import torchvision.transforms as T
import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
def segment_with_deeplab(image_path, save_path='deeplab_result.jpg'):
    # 1. 加载预训练的DeepLabv3
    model = torch.hub.load('pytorch/vision', 'deeplabv3_resnet101', pretrained=True)
    model.eval()
    # 2. 图像预处理
    transform = T.Compose([
        T.Resize(256),
        T.CenterCrop(224),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225])
    ])
    # 3. 读取并处理图像
    image = Image.open(image_path).convert('RGB')
    input_tensor = transform(image).unsqueeze(0)
    # 4. 推理
    with torch.no_grad():
        output = model(input_tensor)['out'][0]
        output_predictions = output.argmax(0).byte().cpu().numpy()
    # 5. 映射到原图尺寸
    output_resized = cv2.resize(output_predictions, (image.size[0], image.size[1]), 
                                interpolation=cv2.INTER_NEAREST)
    # 6. 可视化（PASCAL VOC 21类）
    # 创建一个简单的颜色映射
    palette = np.array([
        [0, 0, 0],        # 背景
        [128, 0, 0],      # 飞机
        [0, 128, 0],      # 自行车
        # ... 更多颜色可根据需要扩展
    ])
    # 这里只显示后21种颜色（PASCAL VOC常用）
    color_map = np.random.randint(0, 255, (21, 3), dtype=np.uint8)
    color_map[0] = [0, 0, 0]  # 背景为黑色
    segmented = color_map[output_resized]
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title("原始图像")
    plt.axis('off')
    plt.subplot(1, 2, 2)
    plt.imshow(segmented)
    plt.title("DeepLabv3分割结果")
    plt.axis('off')
    plt.savefig(save_path)
    plt.show()
    return output_resized
# 使用示例
segment_with_deeplab("/path/to/your/image.jpg")

基于 Contour 的传统图像分割（无深度学习）

如果场景简单（如分割前景和背景），可用OpenCV的轮廓检测方法。

import cv2
import numpy as np
import matplotlib.pyplot as plt
def simple_scene_segmentation(image_path, save_path='simple_seg.jpg'):
    # 1. 读取图像
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 2. 简单方法：使用自适应阈值或Canny边缘检测
    # 方法A：Otsu阈值
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # 方法B：Canny边缘检测 + 形态学操作（推荐）
    edges = cv2.Canny(gray, 50, 150)
    kernel = np.ones((5,5), np.uint8)
    dilated = cv2.dilate(edges, kernel, iterations=2)
    # 3. 找到轮廓
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 4. 创建空白的掩码图像
    mask = np.zeros_like(gray)
    # 5. 绘制所有轮廓（或者面积最大的几个）
    for contour in contours:
        area = cv2.contourArea(contour)
        if area > 500:  # 过滤掉小区域
            cv2.drawContours(mask, [contour], -1, (255), thickness=cv2.FILLED)
    # 6. 可视化结果
    plt.figure(figsize=(15, 5))
    plt.subplot(1, 3, 1)
    plt.imshow(image_rgb)
    plt.title("原始图像")
    plt.subplot(1, 3, 2)
    plt.imshow(binary, cmap='gray')
    plt.title("Otsu二值化")
    plt.subplot(1, 3, 3)
    plt.imshow(mask, cmap='gray')
    plt.title("轮廓分割结果")
    plt.savefig(save_path)
    plt.show()
    return mask
# 使用示例
simple_scene_segmentation("/path/to/your/image.jpg")

总结与选择建议

方法	适用场景	精度	速度	需要GPU
SegFormer	通用语义分割	高	中等	推荐
SAM	万物分割（无需预定义类别）	最高	慢	必须
DeepLabv3	实时分割、移动端	中	快	可选
传统方法	简单场景、前景背景分离	低	极快	不需要