别再手动看水色了！用Python+SVM做个水质图像识别小工具（附完整代码）

用Python打造智能水质检测仪：从图像特征提取到SVM模型部署实战

水产养殖业中，水质监测是确保生物健康生长的关键环节。传统依赖专家肉眼观察水色的方法不仅效率低下，还存在主观性强、难以标准化的问题。本文将带你开发一个基于Python的智能水质检测工具，只需拍摄水体照片，就能自动完成水质分级评估。

1. 项目架构与核心原理

水质图像识别系统的核心在于将颜色特征转化为可量化的数据指标。我们采用颜色矩作为特征提取方法，结合支持向量机(SVM)构建分类模型。整个流程可分为四个关键阶段：

1. 图像采集与预处理：统一图像尺寸，消除光照干扰
2. 特征工程：提取RGB三通道的9个颜色矩特征
3. 模型训练：使用SVM学习不同水色特征
4. 应用部署：将训练好的模型封装为可调用接口

颜色矩是一种紧凑且有效的图像特征表示方法，包含三个层次：

• 一阶颜色矩（均值）：反映图像的整体色调
• 二阶颜色矩（标准差）：描述颜色的分布范围
• 三阶颜色矩（偏度）：体现颜色分布的不对称性

2. 开发环境搭建与数据准备

2.1 工具链配置

推荐使用Python 3.8+环境，主要依赖库包括：

pip install pillow numpy scikit-learn matplotlib opencv-python

核心库功能说明：

2.2 数据集构建

典型的水质图像分类数据集应包含多种水色样本。建议采集以下五类常见水色：

1. 浅绿色（优质水）：51张
2. 灰蓝色（次优水）：44张
3. 黄褐色（轻度污染）：78张
4. 茶褐色（中度污染）：24张
5. 深绿色（重度污染）：6张

注意：样本量不足时可使用数据增强技术，如旋转、添加噪声等，但要确保不改变原始水色特征。

图像命名规范建议采用类别_序号.jpg格式，例如：

1_001.jpg # 浅绿色 3_045.jpg # 黄褐色

3. 图像处理与特征提取实战

3.1 智能图像预处理

from PIL import Image import numpy as np def preprocess_image(image_path, target_size=(256, 256)): """标准化图像尺寸并增强对比度""" img = Image.open(image_path) # 统一尺寸 img = img.resize(target_size) # 转换为RGB模式（去除Alpha通道） if img.mode != 'RGB': img = img.convert('RGB') # 提取中心区域（减少边缘干扰） width, height = img.size left = (width - 200)/2 top = (height - 200)/2 right = (width + 200)/2 bottom = (height + 200)/2 img = img.crop((left, top, right, bottom)) return img

3.2 颜色矩特征计算

颜色矩的计算公式如下：

1. 一阶矩（均值）： $$ \mu_i = \frac{1}{N}\sum_{j=1}^{N}p_{ij} $$

2. 二阶矩（标准差）： $$ \sigma_i = \sqrt{\frac{1}{N}\sum_{j=1}^{N}(p_{ij} - \mu_i)^2} $$

3. 三阶矩（偏度）： $$ s_i = \sqrt[3]{\frac{1}{N}\sum_{j=1}^{N}(p_{ij} - \mu_i)^3} $$

Python实现代码：

def calculate_color_moments(image): """计算RGB三通道的9个颜色矩特征""" # 分离通道 r, g, b = image.split() # 转换为numpy数组 r_array = np.array(r)/255.0 g_array = np.array(g)/255.0 b_array = np.array(b)/255.0 def get_moments(channel): """计算单个通道的三阶颜色矩""" mean = np.mean(channel) std = np.std(channel) skewness = np.mean(np.abs(channel - mean)**3)**(1/3) return [mean, std, skewness] features = [] for channel in [r_array, g_array, b_array]: features.extend(get_moments(channel)) return np.array(features)

4. SVM模型训练与优化

4.1 数据准备与划分

from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler def prepare_dataset(image_dir): """构建特征矩阵和标签向量""" image_files = [f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.png'))] X = np.zeros((len(image_files), 9)) # 9个特征 y = np.zeros(len(image_files)) # 标签 for i, filename in enumerate(image_files): img_path = os.path.join(image_dir, filename) img = preprocess_image(img_path) X[i] = calculate_color_moments(img) # 从文件名提取类别标签（格式：类别_序号.jpg） y[i] = int(filename.split('_')[0]) # 特征标准化 scaler = StandardScaler() X = scaler.fit_transform(X) return train_test_split(X, y, test_size=0.2, random_state=42)

4.2 模型训练与评估

from sklearn.svm import SVC from sklearn.metrics import classification_report def train_evaluate_model(X_train, X_test, y_train, y_test): """训练SVM模型并评估性能""" # 使用类别权重平衡样本不均衡问题 model = SVC(kernel='rbf', class_weight='balanced', gamma='scale') model.fit(X_train, y_train) # 评估指标 train_acc = model.score(X_train, y_train) test_acc = model.score(X_test, y_test) print(f"训练集准确率: {train_acc:.2%}") print(f"测试集准确率: {test_acc:.2%}") print("\n分类报告:") print(classification_report(y_test, model.predict(X_test))) return model

4.3 模型优化技巧

针对小样本数据集的优化策略：

1. 数据增强：

   • 随机旋转（±10°）
   • 添加高斯噪声（σ=0.01）
   • 亮度微调（±10%）

2. 模型参数调优：

   from sklearn.model_selection import GridSearchCV param_grid = { 'C': [0.1, 1, 10], 'gamma': ['scale', 'auto', 0.1, 1], 'kernel': ['rbf', 'poly'] } grid_search = GridSearchCV(SVC(class_weight='balanced'), param_grid, cv=3) grid_search.fit(X_train, y_train) best_model = grid_search.best_estimator_

3. 特征选择：

   • 使用随机森林评估特征重要性
   • 保留贡献度最高的6-7个特征

5. 应用部署与实战技巧

5.1 模型保存与加载

import joblib # 保存模型和标准化器 joblib.dump(model, 'water_quality_model.pkl') joblib.dump(scaler, 'feature_scaler.pkl') # 加载使用 model = joblib.load('water_quality_model.pkl') scaler = joblib.load('feature_scaler.pkl')

5.2 实时预测接口

def predict_water_quality(image_path, model, scaler): """对单张图像进行水质预测""" try: img = preprocess_image(image_path) features = calculate_color_moments(img) scaled_features = scaler.transform([features]) pred_class = int(model.predict(scaled_features)[0]) class_names = { 1: "优质水（浅绿色）", 2: "次优水（灰蓝色）", 3: "轻度污染（黄褐色）", 4: "中度污染（茶褐色）", 5: "重度污染（深绿色）" } return class_names.get(pred_class, "未知水质") except Exception as e: return f"预测错误: {str(e)}"

5.3 移动端集成方案

1. 方案一：Flask REST API

   from flask import Flask, request, jsonify app = Flask(__name__) @app.route('/predict', methods=['POST']) def predict(): if 'file' not in request.files: return jsonify({'error': 'No file uploaded'}) file = request.files['file'] temp_path = f"temp_{file.filename}" file.save(temp_path) result = predict_water_quality(temp_path, model, scaler) os.remove(temp_path) return jsonify({'result': result}) if __name__ == '__main__': app.run(host='0.0.0.0', port=5000)

2. 方案二：PyInstaller打包

   pyinstaller --onefile --add-data 'model.pkl;.' water_quality_app.py

5.4 实际应用中的问题解决

光照条件影响：

• 解决方法：使用白平衡校正

  import cv2 def white_balance(img): result = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) lab = cv2.cvtColor(result, cv2.COLOR_BGR2LAB) l, a, b = cv2.split(lab) clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8)) cl = clahe.apply(l) limg = cv2.merge((cl,a,b)) return Image.fromarray(cv2.cvtColor(limg, cv2.COLOR_LAB2RGB))

拍摄角度问题：

• 建议：在水体中心位置垂直向下拍摄
• 自动校正代码：

  def perspective_correction(img): """透视变换校正""" img_array = np.array(img) gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY) _, thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY) contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) largest = max(contours, key=cv2.contourArea) rect = cv2.minAreaRect(largest) box = cv2.boxPoints(rect) box = np.int0(box) width, height = img.size dst_pts = np.array([[0, height-1], [0, 0], [width-1, 0], [width-1, height-1]], dtype="float32") M = cv2.getPerspectiveTransform(box.astype("float32"), dst_pts) warped = cv2.warpPerspective(img_array, M, (width, height)) return Image.fromarray(warped)