开发轻量级分类模型或作为更大管道检测系统的前置判断模块 排水管道缺陷自动分类 排水管分类数据集 污水管分类数据集的管暗接数据集和沉积、结垢数据集进行分类任务

---

排水污水管道缺陷分类数据集信息表

子数据集名称	类别	图像数量（每类）	总图像数量	说明
支管暗接数据集	缺陷（AJ - 支管暗接）	1,000 张	2,000 张	包含缺陷图像与正常图像各1,000张，用于二分类任务
	正常（无缺陷）	1,000 张		作为对照样本，用于模型区分能力训练
沉积、结垢数据集	沉积（CJ）	40 张	40 张	小样本数据集，展示典型沉积与结垢缺陷（见第6–8张例图）
	结垢（JG）	40 张		图像数量较少，适用于可视化或迁移学习参考

✅ 备注：

• “✔” 表示该子数据集已包含在总数据集中。
• 例图说明：
  • 前5张例图：展示“支管暗接”缺陷与正常状态对比。
  • 第6–8张例图：展示“沉积”与“结垢”典型样本。

---

数据集特点总结

特点	说明
✅ 分类任务导向	专为图像分类设计，适用于二分类或多类别识别模型训练
✅ 标注清晰	每张图像明确标注为“缺陷”或“正常”，类别定义准确
✅ 双场景覆盖	包含常见功能性缺陷：支管暗接、沉积、结垢
✅ 适用性强	可用于卷积神经网络（CNN）、ResNet、EfficientNet 等分类模型训练
⚠️ 注意	“沉积、结垢数据集”图像数量较少（共80张），建议用于测试、演示或结合数据增强使用

---

典型应用方向

• 排水管道缺陷自动分类系统
• 工业视觉质检平台
• AI辅助管道巡检诊断工具
• 计算机视觉教学与实验数据集

---

该数据集适合用于开发轻量级分类模型或作为更大管道检测系统的前置判断模块。

数据集文字说明：

排水污水管道缺陷分类数据集，分为两个子数据集：

✔支管暗接数据集–>缺陷和正常图像各1000张，共2000张。见前5张例图。

✔支管暗接数据集–>缺陷和正常图像各1000张，共2000张。见前5张例图。

✔沉积，结垢数据集–>40张图像，见第6–8张例图。

1

2

3

4
5

✔沉积，结垢数据集–>40张图像，见第6–8张例图。

6
7

8

处理两个子数据集：支管暗接数据集和沉积、结垢数据集，并进行分类任务。以下是详细的步骤和代码示例,仅供参考学习使用。

一、环境准备

1. 安装必要的库

确保同学你已经安装了所需的 Python 库：

pip install numpy matplotlib opencv-python tqdm

二、数据集组织与预处理

假设你的数据集目录结构如下：

drainage_defect_dataset/
├── branch_connection/
│   ├── defect/
│   │   ├── image_0001.jpg
│   │   └── ...
│   ├── normal/
│   │   ├── image_0001.jpg
│   │   └── ...
├── sediment_scaling/
│   ├── image_0001.jpg
│   └── ...

三、数据加载与预处理

1. 加载图像并进行预处理

import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split

def load_images_from_folder(folder):
    images = []
    labels = []
    for filename in os.listdir(folder):
        img_path = os.path.join(folder, filename)
        img = cv2.imread(img_path)
        if img is not None:
            images.append(img)
            labels.append(1 if 'defect' in folder else 0)  # 1 for defect, 0 for normal
    return images, labels

# Load branch connection dataset
branch_defect_folder = 'drainage_defect_dataset/branch_connection/defect'
branch_normal_folder = 'drainage_defect_dataset/branch_connection/normal'

branch_defect_images, branch_defect_labels = load_images_from_folder(branch_defect_folder)
branch_normal_images, branch_normal_labels = load_images_from_folder(branch_normal_folder)

# Combine branch connection data
branch_images = branch_defect_images + branch_normal_images
branch_labels = branch_defect_labels + branch_normal_labels

# Load sediment scaling dataset
sediment_scaling_folder = 'drainage_defect_dataset/sediment_scaling'
sediment_scaling_images, sediment_scaling_labels = load_images_from_folder(sediment_scaling_folder)

# Split data into training and testing sets
branch_train_images, branch_test_images, branch_train_labels, branch_test_labels = train_test_split(
    branch_images, branch_labels, test_size=0.2, random_state=42
)

sediment_scaling_train_images, sediment_scaling_test_images, sediment_scaling_train_labels, sediment_scaling_test_labels = train_test_split(
    sediment_scaling_images, sediment_scaling_labels, test_size=0.2, random_state=42
)

四、模型训练

1. 使用 PyTorch 进行分类模型训练

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

class CustomDataset(Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

# Define a simple CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(16 * 56 * 56, 128)
        self.fc2 = nn.Linear(128, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = x.view(-1, 16 * 56 * 56)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Data transformations
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

# Create datasets and dataloaders
branch_train_dataset = CustomDataset(branch_train_images, branch_train_labels, transform=transform)
branch_test_dataset = CustomDataset(branch_test_images, branch_test_labels, transform=transform)

sediment_scaling_train_dataset = CustomDataset(sediment_scaling_train_images, sediment_scaling_train_labels, transform=transform)
sediment_scaling_test_dataset = CustomDataset(sediment_scaling_test_images, sediment_scaling_test_labels, transform=transform)

branch_train_loader = DataLoader(branch_train_dataset, batch_size=32, shuffle=True)
branch_test_loader = DataLoader(branch_test_dataset, batch_size=32, shuffle=False)

sediment_scaling_train_loader = DataLoader(sediment_scaling_train_dataset, batch_size=32, shuffle=True)
sediment_scaling_test_loader = DataLoader(sediment_scaling_test_dataset, batch_size=32, shuffle=False)

# Initialize the model, loss function, and optimizer
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in branch_train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {running_loss / len(branch_train_loader)}")

# Evaluation
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in branch_test_loader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f"Accuracy: {100 * correct / total}%")

五、推理与评估

1. 模型推理

def predict_image(model, image_path):
    image = cv2.imread(image_path)
    image = transform(image)
    image = image.unsqueeze(0)
    output = model(image)
    _, predicted = torch.max(output.data, 1)
    return predicted.item()

test_image_path = 'path_to_test_image.jpg'
prediction = predict_image(model, test_image_path)
print(f"Prediction: {'Defect' if prediction == 1 else 'Normal'}")

六、总结

以上代码展示了如何处理排水污水管道缺陷分类数据集，包括数据加载、预处理、模型训练、推理和评估。