Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

第三十八 计算机视觉(CV, Computer Vision)

第三十八 计算机视觉
├── 什么是计算机视觉(CV)
├── 图像基础
│   ├── 像素、通道、图像表示
│   └── 图像预处理
├── 卷积运算
│   ├── 卷积的定义与数学表达
│   └── 卷积与图像处理
├── 卷积神经网络(CNN)
│   ├── 卷积层、池化层、全连接层
│   └── 经典模型(LeNet、VGG、ResNet)
├── 目标检测与图像分割
│   ├── 目标检测(YOLO、R-CNN)
│   └── 图像分割(语义分割、实例分割)
├── Rust CV生态
│   ├── image crate 图像处理
│   ├── candle 深度学习推理
│   └── 其他 CV 相关 crate
└── 实战:图像分类(CIFAR-10)
    └── 用Rust + candle实现

计算机视觉(Computer Vision, CV)是人工智能领域中最具感知力的分支之一。人类的视觉系统可以毫不费力地识别物体、理解场景、判断距离,而计算机视觉的目标正是让机器“看懂“世界——从像素矩阵中提取有意义的语义信息。

$$\text{图像(像素矩阵)} \xrightarrow{\text{计算机视觉}} \text{语义理解(分类、检测、分割)}$$


一、什么是计算机视觉

1.1 定义

计算机视觉是研究如何使机器从数字图像或多维数据中获取高层理解的学科。它涵盖了从低层图像处理(滤波、增强)到高层语义理解(识别、推理)的完整链路。

1.2 CV 的核心任务层次

层次任务示例
低层图像预处理、增强灰度化、去噪、边缘检测
中层特征提取、描述SIFT、HOG、卷积特征
高层语义理解分类、检测、分割、生成

1.3 发展简史

  • 1960s:Block World,积木世界识别
  • 1980s:David Marr 视觉计算理论,边缘检测
  • 1990s:SIFT、HOG 等手工特征
  • 2012:AlexNet 在 ImageNet 竞赛中夺冠,深度学习时代开启
  • 2015+:ResNet、YOLO、U-Net 等经典模型涌现
  • 2020s:Vision Transformer (ViT)、多模态大模型

二、图像基础

2.1 像素与图像表示

一幅数字图像可以表示为一个二维矩阵。对于灰度图像,每个像素是一个标量值;对于彩色图像,每个像素是一个向量。

灰度图像:$H \times W$ 矩阵,像素值 $I(x, y) \in [0, 255]$

$$I_{\text{gray}} = \begin{pmatrix} 128 & 64 & 32 \ 200 & 150 & 100 \ 255 & 0 & 128 \end{pmatrix}$$

彩色图像(RGB):$H \times W \times 3$ 张量,三个通道分别为红(R)、绿(G)、蓝(B)

$$I_{\text{rgb}}(x, y) = \big[ R(x,y),\ G(x,y),\ B(x,y) \big]$$

其中每个通道的值域为 $[0, 255]$,在深度学习中通常归一化到 $[0, 1]$。

2.2 RGB 通道

RGB 三原色通过加色混合可以产生人眼能感知的大部分颜色:

通道颜色纯色值 (R, G, B)
R红色(255, 0, 0)
G绿色(0, 255, 0)
B蓝色(0, 0, 255)
白色(255, 255, 255)
黑色(0, 0, 0)

2.3 用 Rust 表示图像

/// 像素结构体
#[derive(Debug, Clone, Copy)]
struct Pixel {
    r: u8,
    g: u8,
    b: u8,
}

/// 简易图像结构
struct Image {
    width: usize,
    height: usize,
    pixels: Vec<Pixel>,
}

impl Image {
    fn new(width: usize, height: usize) -> Self {
        Image {
            width,
            height,
            pixels: vec![Pixel { r: 0, g: 0, b: 0 }; width * height],
        }
    }

    fn get_pixel(&self, x: usize, y: usize) -> &Pixel {
        &self.pixels[y * self.width + x]
    }

    fn set_pixel(&mut self, x: usize, y: usize, pixel: Pixel) {
        self.pixels[y * self.width + x] = pixel;
    }
}

fn main() {
    let mut img = Image::new(3, 2);
    // 设置一个红色像素
    img.set_pixel(0, 0, Pixel { r: 255, g: 0, b: 0 });
    // 设置一个绿色像素
    img.set_pixel(1, 0, Pixel { r: 0, g: 255, b: 0 });

    println!("图像尺寸: {}x{}", img.width, img.height);
    println!("像素 (0,0): {:?}", img.get_pixel(0, 0));
    println!("像素 (1,0): {:?}", img.get_pixel(1, 0));
}

三、图像预处理

图像预处理是 CV 流水线的第一步,目的是将原始图像转换为适合后续处理的标准化形式。

3.1 缩放(Resize)

缩放是最基本的预处理操作。常见的插值算法包括最近邻插值、双线性插值和双三次插值。

$$I’(x’, y’) = \sum_{i} \sum_{j} I(x_i, y_j) \cdot w(x_i - x’) \cdot w(y_j - y’)$$

其中 $w$ 为插值核函数。

use image::{imageops, DynamicImage, GenericImageView, ImageFormat};

fn main() {
    // 加载图像
    let img = image::open("input.jpg").expect("无法打开图像");
    println!("原始尺寸: {}x{}", img.width(), img.height());

    // 使用不同滤波器缩放
    let methods = [
        ("nearest",  imageops::FilterType::Nearest),
        ("triangle", imageops::FilterType::Triangle),
        ("catmull",  imageops::FilterType::CatmullRom),
        ("gaussian", imageops::FilterType::Gaussian),
        ("lanczos",  imageops::FilterType::Lanczos3),
    ];

    for (name, filter) in &methods {
        let resized = img.resize(224, 224, *filter);
        resized.save(format!("resized_{}.png", name)).unwrap();
        println!("已保存 resized_{}.png (224x224)", name);
    }
}

3.2 灰度化

将彩色图像转换为灰度图像,常用加权平均法(ITU-R BT.601 标准):

$$I_{\text{gray}} = 0.299 \cdot R + 0.587 \cdot G + 0.114 \cdot B$$

人眼对绿色最敏感,对蓝色最不敏感,因此权重不同。

use image::{DynamicImage, GenericImageView, Rgb, Luma};

fn to_grayscale(img: &DynamicImage) -> image::GrayImage {
    let rgb = img.to_rgb8();
    let (width, height) = rgb.dimensions();
    let mut gray = image::GrayImage::new(width, height);

    for y in 0..height {
        for x in 0..width {
            let Rgb([r, g, b]) = rgb.get_pixel(x, y);
            // ITU-R BT.601 加权灰度
            let gray_val = 0.299 * *r as f32
                         + 0.587 * *g as f32
                         + 0.114 * *b as f32;
            gray.put_pixel(x, y, Luma([gray_val as u8]));
        }
    }
    gray
}

fn main() {
    let img = image::open("input.jpg").unwrap();
    let gray = to_grayscale(&img);
    gray.save("grayscale.png").unwrap();
    println!("灰度图像已保存");

    // image crate 也内置了灰度转换
    let gray2 = img.to_luma8();
    gray2.save("grayscale_builtin.png").unwrap();
}

3.3 滤波(卷积滤波)

滤波是图像处理的核心操作,通过卷积核(kernel)与图像卷积实现各种效果。

均值滤波(模糊)

$$K_{\text{mean}} = \frac{1}{9} \begin{pmatrix} 1 & 1 & 1 \ 1 & 1 & 1 \ 1 & 1 & 1 \end{pmatrix}$$

高斯滤波

$$K_{\text{gauss}} = \frac{1}{16} \begin{pmatrix} 1 & 2 & 1 \ 2 & 4 & 2 \ 1 & 2 & 1 \end{pmatrix}$$

Sobel 边缘检测(水平方向):

$$K_{\text{sobel_x}} = \begin{pmatrix} -1 & 0 & 1 \ -2 & 0 & 2 \ -1 & 0 & 1 \end{pmatrix}$$

use image::{GrayImage, GenericImageView, Luma};

/// 手动实现 3x3 卷积
fn convolve3x3(img: &GrayImage, kernel: &[f32; 9]) -> GrayImage {
    let (width, height) = img.dimensions();
    let mut output = GrayImage::new(width, height);

    for y in 1..height - 1 {
        for x in 1..width - 1 {
            let mut sum = 0.0f32;
            let mut ki = 0;
            for ky in -1..=1 {
                for kx in -1..=1 {
                    let px = (x as i32 + kx) as u32;
                    let py = (y as i32 + ky) as u32;
                    let Luma([val]) = img.get_pixel(px, py);
                    sum += val as f32 * kernel[ki];
                    ki += 1;
                }
            }
            output.put_pixel(x, y, Luma([sum.clamp(0.0, 255.0) as u8]));
        }
    }
    output
}

fn main() {
    let img = image::open("input.jpg").unwrap().to_luma8();

    // 均值滤波(模糊)
    let mean_kernel: [f32; 9] = [1.0/9.0; 9];
    let blurred = convolve3x3(&img, &mean_kernel);
    blurred.save("blurred.png").unwrap();

    // 高斯滤波
    let gauss_kernel: [f32; 9] =
        [1.0/16.0, 2.0/16.0, 1.0/16.0,
         2.0/16.0, 4.0/16.0, 2.0/16.0,
         1.0/16.0, 2.0/16.0, 1.0/16.0];
    let gaussian = convolve3x3(&img, &gauss_kernel);
    gaussian.save("gaussian.png").unwrap();

    // Sobel 边缘检测
    let sobel_x: [f32; 9] = [-1.0, 0.0, 1.0, -2.0, 0.0, 2.0, -1.0, 0.0, 1.0];
    let edges = convolve3x3(&img, &sobel_x);
    edges.save("edges.png").unwrap();

    println!("滤波结果已保存");
}

四、卷积运算

4.1 卷积的数学定义

在信号处理中,两个函数 $f$ 和 $g$ 的卷积定义为:

$$(f * g)(t) = \int_{-\infty}^{+\infty} f(\tau) \cdot g(t - \tau) , d\tau$$

在离散情况下:

$$(f * g)[n] = \sum_{m=-\infty}^{+\infty} f[m] \cdot g[n - m]$$

4.2 二维卷积(图像卷积)

对于图像 $I$ 和卷积核 $K$,二维卷积运算为:

$$(I * K)(i, j) = \sum_{m} \sum_{n} I(i + m, j + n) \cdot K(m, n)$$

注意: 在深度学习中,实际使用的是“互相关“(cross-correlation)而非严格意义上的卷积,区别在于核是否翻转。实践中两者效果等价,因为核的参数是学习得到的。

4.3 步长与填充

  • 步长(Stride):卷积核每次移动的像素数,记为 $s$
  • 填充(Padding):在图像边缘补零,记为 $p$
  • 输出尺寸公式

$$O = \left\lfloor \frac{W - K + 2P}{S} \right\rfloor + 1$$

其中 $W$ 为输入尺寸,$K$ 为核大小,$P$ 为填充大小,$S$ 为步长。

/// 二维卷积实现(互相关)
fn conv2d(
    input: &[Vec<f32>],
    kernel: &[Vec<f32>],
    stride: usize,
    padding: usize,
) -> Vec<Vec<f32>> {
    let h_in = input.len();
    let w_in = input[0].len();
    let k_h = kernel.len();
    let k_w = kernel[0].len();

    // 填充输入
    let h_pad = h_in + 2 * padding;
    let w_pad = w_in + 2 * padding;
    let mut padded = vec![vec![0.0f32; w_pad]; h_pad];
    for i in 0..h_in {
        for j in 0..w_in {
            padded[i + padding][j + padding] = input[i][j];
        }
    }

    // 计算输出尺寸
    let h_out = (h_pad - k_h) / stride + 1;
    let w_out = (w_pad - k_w) / stride + 1;
    let mut output = vec![vec![0.0f32; w_out]; h_out];

    for i in 0..h_out {
        for j in 0..w_out {
            let mut sum = 0.0f32;
            for ki in 0..k_h {
                for kj in 0..k_w {
                    let pi = i * stride + ki;
                    let pj = j * stride + kj;
                    sum += padded[pi][pj] * kernel[ki][kj];
                }
            }
            output[i][j] = sum;
        }
    }
    output
}

fn main() {
    // 5x5 输入
    let input: Vec<Vec<f32>> = (0..5)
        .map(|i| (0..5).map(|j| (i * 5 + j) as f32).collect())
        .collect();

    // 3x3 卷积核(边缘检测)
    let kernel: Vec<Vec<f32>> = vec![
        vec![-1.0, -1.0, -1.0],
        vec![-1.0,  8.0, -1.0],
        vec![-1.0, -1.0, -1.0],
    ];

    let output = conv2d(&input, &kernel, 1, 1);
    println!("卷积输出 (3x3):");
    for row in &output {
        println!("  {:?}", row.iter().map(|v| format!("{:6.1}", v)).collect::<Vec<_>>());
    }
}

五、卷积神经网络(CNN)

5.1 为什么需要 CNN

传统图像处理依赖手工设计的特征(如 SIFT、HOG),而 CNN 能够自动从数据中学习特征层次结构:

$$\text{像素} \rightarrow \text{边缘} \rightarrow \text{纹理} \rightarrow \text{部件} \rightarrow \text{物体}$$

CNN 的三大核心特性:

  • 局部连接:每个神经元只与局部区域连接
  • 权值共享:同一卷积核在整幅图像上滑动共享参数
  • 平移不变性:物体出现在图像不同位置,仍能被识别

5.2 卷积层

卷积层是 CNN 的核心,通过多个卷积核提取不同特征:

$$Y^{(l)} = \sigma\left( W^{(l)} * X^{(l-1)} + b^{(l)} \right)$$

其中 $\sigma$ 为激活函数(通常为 ReLU),$W^{(l)}$ 为第 $l$ 层的卷积核权重,$b^{(l)}$ 为偏置。

ReLU 激活函数

$$\text{ReLU}(x) = \max(0, x)$$

ReLU 的优势:计算简单、缓解梯度消失问题。

5.3 池化层

池化层用于降低特征图的空间维度,减少计算量和参数量。

最大池化(Max Pooling)

$$Y(i, j) = \max_{(m,n) \in \mathcal{R}_{ij}} X(m, n)$$

平均池化(Average Pooling)

$$Y(i, j) = \frac{1}{|\mathcal{R}{ij}|} \sum{(m,n) \in \mathcal{R}_{ij}} X(m, n)$$

其中 $\mathcal{R}_{ij}$ 为池化窗口覆盖的区域。

/// 最大池化 2x2, stride=2
fn max_pool2x2(input: &[Vec<f32>]) -> Vec<Vec<f32>> {
    let h = input.len() / 2;
    let w = input[0].len() / 2;
    let mut output = vec![vec![0.0f32; w]; h];

    for i in 0..h {
        for j in 0..w {
            let r = i * 2;
            let c = j * 2;
            output[i][j] = input[r][c]
                .max(input[r][c + 1])
                .max(input[r + 1][c])
                .max(input[r + 1][c + 1]);
        }
    }
    output
}

/// 平均池化 2x2, stride=2
fn avg_pool2x2(input: &[Vec<f32>]) -> Vec<Vec<f32>> {
    let h = input.len() / 2;
    let w = input[0].len() / 2;
    let mut output = vec![vec![0.0f32; w]; h];

    for i in 0..h {
        for j in 0..w {
            let r = i * 2;
            let c = j * 2;
            let sum = input[r][c] + input[r][c + 1]
                    + input[r + 1][c] + input[r + 1][c + 1];
            output[i][j] = sum / 4.0;
        }
    }
    output
}

fn main() {
    let input: Vec<Vec<f32>> = vec![
        vec![1.0, 2.0, 3.0, 4.0],
        vec![5.0, 6.0, 7.0, 8.0],
        vec![9.0, 10.0, 11.0, 12.0],
        vec![13.0, 14.0, 15.0, 16.0],
    ];

    let max_pooled = max_pool2x2(&input);
    println!("最大池化: {:?}", max_pooled);
    // [[6.0, 8.0], [14.0, 16.0]]

    let avg_pooled = avg_pool2x2(&input);
    println!("平均池化: {:?}", avg_pooled);
    // [[3.5, 5.5], [11.5, 13.5]]
}

5.4 全连接层

全连接层将卷积层提取的高层特征展平后映射到输出空间:

$$\mathbf{y} = \sigma(W \cdot \mathbf{x} + \mathbf{b})$$

其中 $\mathbf{x}$ 是展平后的特征向量,$W$ 是权重矩阵,$\mathbf{b}$ 是偏置向量。

对于 $K$ 分类问题,输出层使用 Softmax:

$$\text{Softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}$$

5.5 经典 CNN 模型

LeNet-5(1998)

LeNet-5 是 Yann LeCun 提出的最早的成功 CNN 之一,用于手写数字识别。

类型输出尺寸参数量
1卷积层 (6@5x5)28x28x6156
2池化层 (2x2)14x14x60
3卷积层 (16@5x5)10x10x162,416
4池化层 (2x2)5x5x160
5全连接层 (120)12048,120
6全连接层 (84)8410,164
7输出层 (10)10850

VGG-16(2014)

VGG 的核心思想是使用小卷积核(3x3)和深层堆叠:

$$\text{两个 } 3 \times 3 \text{ 卷积的感受野} = \text{一个 } 5 \times 5 \text{ 卷积的感受野}$$

但参数更少:$2 \times (3 \times 3 \times C^2) = 18C^2 < 25C^2 = 5 \times 5 \times C^2$

VGG-16 共 16 个权重层,约 1.38 亿参数。

ResNet(2015)

ResNet 引入了残差连接(Skip Connection),解决了深层网络的退化问题:

$$\mathcal{F}(x) = \mathcal{H}(x) - x \quad \Rightarrow \quad \mathcal{H}(x) = \mathcal{F}(x) + x$$

网络学习的是残差 $\mathcal{F}(x) = \mathcal{H}(x) - x$,而非直接学习映射 $\mathcal{H}(x)$。

/// 简化的残差块概念演示
fn residual_block(input: &[f32], weights: &[f32], bias: &[f32]) -> Vec<f32> {
    // F(x) = W * x + b (简化为逐元素运算)
    let f_x: Vec<f32> = input
        .iter()
        .zip(weights.iter())
        .zip(bias.iter())
        .map(|((&x, &w), &b)| (x * w + b).max(0.0)) // ReLU
        .collect();

    // H(x) = F(x) + x (残差连接)
    let h_x: Vec<f32> = input.iter().zip(f_x.iter()).map(|(&x, &f)| x + f).collect();

    h_x
}

fn main() {
    let input = vec![1.0, 2.0, 3.0, 4.0, 5.0];
    let weights = vec![0.5, 0.8, 0.3, 0.6, 0.4];
    let bias = vec![0.1, -0.1, 0.2, -0.2, 0.1];

    let output = residual_block(&input, &weights, &bias);
    println!("残差块输出: {:?}", output);
    // [1.6, 2.5, 2.9, 4.4, 5.3]
}

六、目标检测

6.1 目标检测概述

目标检测不仅要识别图像中有什么物体,还要定位它们的位置。输出为边界框(Bounding Box)和类别标签。

$$\text{目标检测输出} = {(x_{\min}, y_{\min}, x_{\max}, y_{\max}, \text{class}, \text{confidence})}$$

6.2 两阶段检测器:R-CNN 系列

R-CNN(Region-based CNN)系列采用“先选区域,再分类“的策略:

  1. R-CNN:Selective Search 生成候选区域 -> CNN 提取特征 -> SVM 分类
  2. Fast R-CNN:引入 RoI Pooling,共享卷积特征
  3. Faster R-CNN:用 RPN(Region Proposal Network)替代 Selective Search

6.3 单阶段检测器:YOLO

YOLO(You Only Look Once)将检测视为回归问题,一次前向传播完成检测:

$$\text{YOLO 输出} = S \times S \times (B \times 5 + C)$$

其中 $S \times S$ 为网格数,$B$ 为每个网格的边界框数,$C$ 为类别数。

YOLO 的优势:速度快,适合实时检测。

模型速度 (FPS)mAP特点
R-CNN~0.0758.4精度高,速度慢
Faster R-CNN~773.2精度与速度平衡
YOLOv3~4557.9实时检测
YOLOv8~100+53.9最新架构,多任务

七、图像分割

7.1 语义分割

语义分割为图像中的每个像素分配一个类别标签,不区分同类的不同实例。

$$f: H \times W \rightarrow {1, 2, \ldots, K}$$

经典模型:FCN(全卷积网络)、U-Net。

U-Net 采用编码器-解码器结构,通过跳跃连接(skip connection)保留空间细节:

编码器(下采样)          解码器(上采样)
64  -> 128  -> 256  -> 512  -> 1024
 |      |      |      |        |
 +------+------+------+--------+ (跳跃连接)

7.2 实例分割

实例分割不仅为每个像素分类,还要区分同类的不同实例。

$$f: H \times W \rightarrow {(c_i, m_i)}_{i=1}^{N}$$

其中 $c_i$ 为类别,$m_i$ 为实例掩码。

经典模型:Mask R-CNN。

7.3 分割与检测对比

任务输出粒度是否区分实例典型模型
图像分类整幅图像ResNet
目标检测边界框YOLO
语义分割像素级U-Net
实例分割像素级掩码Mask R-CNN

八、Rust CV 生态

8.1 image crate

image 是 Rust 生态中最核心的图像处理库,支持多种格式的编解码和基本图像操作。

# Cargo.toml
[dependencies]
image = "0.25"
use image::{DynamicImage, GenericImageView, ImageFormat, imageops};

fn main() {
    // 1. 读取图像(支持 JPEG, PNG, WebP, GIF, BMP, TIFF 等)
    let img = image::open("photo.jpg").expect("打开图像失败");
    println!("尺寸: {}x{}", img.width(), img.height());

    // 2. 裁剪
    let cropped = img.crop_imm(100, 100, 300, 300);
    cropped.save("cropped.png").unwrap();

    // 3. 旋转
    let rotated = img.rotate90();
    rotated.save("rotated.png").unwrap();

    // 4. 翻转
    let flipped = imageops::flip_horizontal(&img);
    flipped.save("flipped_h.png").unwrap();

    // 5. 调整亮度/对比度
    let adjusted = img.adjust_contrast(1.5);
    adjusted.save("contrast.png").unwrap();

    // 6. 格式转换
    let rgb = img.to_rgb8();
    let gray = img.to_luma8();

    // 7. 保存为不同格式
    img.save("output.webp").unwrap(); // 自动推断格式
    let mut out = std::fs::File::create("output.bmp").unwrap();
    img.write_to(&mut out, ImageFormat::Bmp).unwrap();

    println!("图像处理完成");
}

8.2 candle 深度学习推理

candle 是 HuggingFace 开发的 Rust 深度学习框架,支持模型推理和训练。

# Cargo.toml
[dependencies]
candle-core = "0.8"
candle-nn = "0.8"
use candle_core::{Tensor, Device};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    // 创建一个 4 维张量,模拟一批图像 (batch, channels, height, width)
    // 形状: [2, 3, 224, 224],模拟 2 张 RGB 图像
    let images = Tensor::randn(0f32, 1.0, (2, 3, 224, 224), &device)?;

    println!("输入张量形状: {:?}", images.shape());

    // 模拟卷积层权重: 16 个 3x3 卷积核,输入 3 通道
    let weight = Tensor::randn(0f32, 0.02, (16, 3, 3, 3), &device)?;
    let bias = Tensor::zeros(16, &device)?;

    // 执行卷积运算 (padding=1)
    let conv = images.conv2d(&weight, &bias, 1, 1)?;
    println!("卷积后形状: {:?}", conv.shape()); // [2, 16, 224, 224]

    // ReLU 激活
    let activated = conv.relu()?;
    println!("ReLU 后形状: {:?}", activated.shape());

    // 最大池化 2x2
    let pooled = activated.max_pool2d(2)?;
    println!("池化后形状: {:?}", pooled.shape()); // [2, 16, 112, 112]

    Ok(())
}

8.3 其他 CV 相关 crate

crate功能说明
image图像编解码与处理Rust CV 基础库
candle深度学习框架HuggingFace 出品,支持推理与训练
burn深度学习框架纯 Rust,支持多后端
nalgebra线性代数矩阵运算、变换
imageproc图像处理算法形态学、阈值分割等
opencv-rustOpenCV 绑定完整的 OpenCV 功能
tractONNX/TFLite 推理轻量级推理引擎
tchPyTorch 绑定C++ torch 的 Rust 封装
qr2termQR 码终端显示将 QR 码渲染到终端

8.4 用 candle 加载预训练模型进行图像分类

use candle_core::{Device, Tensor};
use candle_nn::{VarMap, Module, Conv2d, Conv2dConfig, Linear};
use candle_nn::init;

/// 构建简易 CNN 模型(用于演示,非预训练)
fn build_simple_cnn(vs: &candle_nn::VarBuilder) -> Result<Box<dyn Module>, Box<dyn std::error::Error>> {
    // Conv1: 3 -> 32, 3x3
    let conv1 = candle_nn::conv2d(3, 32, 3, Conv2dConfig::default(), vs.pp("conv1"))?;
    // Conv2: 32 -> 64, 3x3
    let conv2 = candle_nn::conv2d(32, 64, 3, Conv2dConfig::default(), vs.pp("conv2"))?;
    // FC: 64*6*6 -> 128
    let fc1 = candle_nn::linear(64 * 6 * 6, 128, vs.pp("fc1"))?;
    // FC: 128 -> 10 (CIFAR-10 的 10 个类别)
    let fc2 = candle_nn::linear(128, 10, vs.pp("fc2"))?;

    Ok(Box::new(move |xs: &Tensor| -> Result<Tensor, candle_core::Error> {
        let xs = xs.apply(&conv1)?.relu()?.max_pool2d(2)?;
        let xs = xs.apply(&conv2)?.relu()?.max_pool2d(2)?;
        let xs = xs.flatten_from(1)?;
        let xs = xs.apply(&fc1)?.relu()?;
        xs.apply(&fc2)
    }))
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;
    let mut varmap = VarMap::new();
    let vs = candle_nn::VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);

    let model = build_simple_cnn(&vs)?;

    // 模拟输入: batch=1, channels=3, height=32, width=32 (CIFAR-10 尺寸)
    let input = Tensor::randn(0f32, 1.0, (1, 3, 32, 32), &device)?;
    let output = model.forward(&input)?;

    println!("模型输出形状: {:?}", output.shape()); // [1, 10]
    println!("原始输出: {}", output.to_vec2::<f32>()?);

    // Softmax 获取概率分布
    let probs = candle_nn::ops::softmax(&output, 1)?;
    println!("概率分布: {:?}", probs.to_vec2::<f32>()?);

    // 获取预测类别
    let pred = probs.argmax(1)?;
    println!("预测类别: {:?}", pred.to_vec1::<u32>()?);

    Ok(())
}

九、实战:图像分类(CIFAR-10)

CIFAR-10 是经典的图像分类数据集,包含 10 个类别、60000 张 32x32 彩色图像。

9.1 数据集概览

类别示例内容
airplane飞机
automobile汽车
bird
cat
deer鹿
dog
frog
horse
ship
truck卡车

9.2 完整训练流程

use candle_core::{Device, Tensor, DType};
use candle_nn::{VarMap, Module, Conv2d, Conv2dConfig, Linear, AdamW, Optimizer, loss};
use candle_datasets::cifar10;

/// 简易 CNN 用于 CIFAR-10 分类
struct CifarNet {
    conv1: Conv2d,
    conv2: Conv2d,
    conv3: Conv2d,
    fc1: Linear,
    fc2: Linear,
}

impl CifarNet {
    fn new(vs: candle_nn::VarBuilder) -> Result<Self, candle_core::Error> {
        let conv1 = candle_nn::conv2d(3, 32, 3, Conv2dConfig::with_padding(1), vs.pp("c1"))?;
        let conv2 = candle_nn::conv2d(32, 64, 3, Conv2dConfig::with_padding(1), vs.pp("c2"))?;
        let conv3 = candle_nn::conv2d(64, 64, 3, Conv2dConfig::with_padding(1), vs.pp("c3"))?;
        let fc1 = candle_nn::linear(64 * 4 * 4, 64, vs.pp("fc1"))?;
        let fc2 = candle_nn::linear(64, 10, vs.pp("fc2"))?;
        Ok(Self { conv1, conv2, conv3, fc1, fc2 })
    }

    fn forward(&self, xs: &Tensor) -> Result<Tensor, candle_core::Error> {
        let xs = xs.apply(&self.conv1)?.relu()?.max_pool2d(2)?;  // [B, 32, 16, 16]
        let xs = xs.apply(&self.conv2)?.relu()?.max_pool2d(2)?;  // [B, 64, 8, 8]
        let xs = xs.apply(&self.conv3)?.relu()?.max_pool2d(2)?;  // [B, 64, 4, 4]
        let xs = xs.flatten_from(1)?;                            // [B, 1024]
        let xs = xs.apply(&self.fc1)?.relu()?;
        xs.apply(&self.fc2)
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    // 加载 CIFAR-10 数据集
    let (train_images, train_labels) = cifar10::load_train()?;
    let (test_images, test_labels) = cifar10::load_test()?;

    println!("训练集: {} 张图像", train_images.len());
    println!("测试集: {} 张图像", test_images.len());

    // 构建模型
    let mut varmap = VarMap::new();
    let vs = candle_nn::VarBuilder::from_varmap(&varmap, DType::F32, &device);
    let model = CifarNet::new(vs)?;

    // 优化器
    let mut optimizer = candle_nn::AdamW::new(&varmap, 0.001f64)?;

    // 训练循环
    let epochs = 10;
    let batch_size = 64;

    for epoch in 1..=epochs {
        let mut total_loss = 0.0f32;
        let mut correct = 0usize;
        let mut total = 0usize;

        // 遍历 mini-batch
        for i in (0..train_images.len()).step_by(batch_size) {
            let end = (i + batch_size).min(train_images.len());
            let batch_images: Vec<Tensor> = train_images[i..end]
                .iter()
                .map(|img| Tensor::from_data(img.clone(), &device).unwrap())
                .collect();
            let batch_labels = Tensor::from_vec(
                train_labels[i..end].to_vec(),
                (end - i,) as usize,
                &device,
            )?;

            let images = Tensor::stack(&batch_images, 0)?;
            let logits = model.forward(&images)?;
            let loss = loss::cross_entropy(&logits, &batch_labels)?;

            optimizer.backward_step(&loss)?;
            total_loss += loss.to_scalar::<f32>()?;

            // 计算准确率
            let preds = logits.argmax(1)?;
            let labels = batch_labels;
            correct += preds.eq(&labels)?.to_scalar::<u32>()? as usize;
            total += end - i;
        }

        println!(
            "Epoch {}/{} - Loss: {:.4} - Acc: {:.2}%",
            epoch, epochs,
            total_loss / (train_images.len() / batch_size) as f32,
            100.0 * correct as f32 / total as f32
        );
    }

    Ok(())
}

十、总结

核心知识回顾

主题核心要点
图像基础像素、RGB 通道、$H \times W \times C$ 张量表示
图像预处理缩放(插值)、灰度化(加权平均)、滤波(卷积核)
卷积运算二维卷积、步长、填充、输出尺寸公式
CNN 核心卷积层(特征提取)、池化层(降维)、全连接层(分类)
经典模型LeNet-5(奠基)、VGG(深层堆叠)、ResNet(残差连接)
目标检测R-CNN(两阶段)、YOLO(单阶段实时检测)
图像分割语义分割(FCN/U-Net)、实例分割(Mask R-CNN)
Rust 生态image(图像处理)、candle(深度学习)、tract(推理)

CV 任务层次总览

图像分类 ──── "这是什么?"
    │
目标检测 ──── "在哪里?是什么?"(边界框 + 类别)
    │
语义分割 ──── "每个像素属于什么类别?"
    │
实例分割 ──── "每个像素属于哪个实例?"

练习建议

  1. 基础练习:使用 image crate 读取一张彩色图片,分别提取 R、G、B 三个通道并保存为灰度图。
  2. 卷积实现:手动实现 5x5 高斯卷积核,对灰度图像进行平滑处理,并与 imageproc 库的结果对比。
  3. CNN 构建:使用 candle 构建一个包含 3 个卷积层 + 2 个全连接层的 CNN,在 CIFAR-10 上训练并记录准确率。
  4. 模型对比:分别实现 LeNet 和 VGG 风格的网络,比较在相同数据集上的参数量和训练效果。
  5. 进阶挑战:使用 tract crate 加载一个 ONNX 格式的预训练 ResNet 模型,对本地图片进行推理分类。
  6. 边缘检测:实现 Sobel 算子(水平和垂直方向),计算梯度幅值,观察不同图像的边缘检测效果。
  7. 图像分割:基于阈值分割实现一个简单的二值化分割器,将前景与背景分离。