AI模型训练猫狗皮肤病医疗数据集训练用于猫狗皮肤病诊断的AI模型

本文链接：https://blog.csdn.net/2401_88440984/article/details/148046208

如何用于AI模型训练

文章目录

医疗数据的猫狗皮肤病医疗数据集AI模型训练
di:03
10,000 只宠物
– 7 种宠物疾病和 4 种猫皮肤病的图像
500,000 张宠物（狗和猫）图片
250,000 张疾病图像
如何用于AI模型训练
在这里插入图片描述

训练一套用于猫狗皮肤病诊断的AI模型涉及多个步骤，包括数据准备、环境搭建、模型选择与配置、训练过程以及性能评估。
示例代码，

1. 数据准备

文件结构

确保你的数据集文件结构如下：

pet_skin_disease_dataset/
├── images/
│   ├── dog/
│   └── cat/
└── labels/
    ├── dog/
    └── cat/

每张图像对应一个标签文件（可以是分类标签或边界框标注），根据你的具体需求进行调整。

类别映射文件（classes.txt）

创建一个 classes.txt 文件，列出所有疾病类别。例如：

disease1
disease2
...
cat_skin_disease1
cat_skin_disease2
...

2. 环境搭建

安装必要的依赖包：

pip install torch torchvision torchaudio
git clone https://github.com/pytorch/vision.git
cd vision
pip install -e .

3. 数据预处理

创建自定义Dataset类以加载和预处理图像数据。

import os
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms

class PetSkinDiseaseDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        self.images = []
        self.labels = []

        for label in os.listdir(img_dir):
            path = os.path.join(img_dir, label)
            if not os.path.isdir(path):
                continue
            for img_name in os.listdir(path):
                if img_name.endswith('.jpg') or img_name.endswith('.png'):
                    self.images.append(os.path.join(path, img_name))
                    self.labels.append(label)

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

    def __getitem__(self, idx):
        img_path = self.images[idx]
        image = Image.open(img_path).convert("RGB")
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label

4. 模型选择与配置

我们可以使用ResNet作为基础模型，并对其进行微调以适应我们的任务。

import torchvision.models as models
import torch.nn as nn

model = models.resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(classes))  # classes为你的疾病类别数

5. 训练过程

数据加载器

transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整尺寸以适应模型输入要求
    transforms.ToTensor(),
])

dataset = PetSkinDiseaseDataset(img_dir='path/to/pet_skin_disease_dataset/images', transform=transform)
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

训练循环

import torch.optim as optim
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item() * inputs.size(0)
    
    print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader.dataset)}")

6. 性能评估

在验证集上评估模型性能：

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

print(f'Accuracy of the model on the validation images: {100 * correct / total}%')