如何用yolov5训练使用——滑坡数据集 自然灾害洪水滑坡检测YOLO数据集模型2339张 2类 ，按照8比2划分为训练集和验证集，其中训练集【1871】张，验证集【468】张，模型分为【2】类】

    def __len__(self):
    return len(self.image_files)
    
    def __getitem__(self, index):
    img_path = os.path.join(self.image_dir, self.image_files[index])
    label_path = os.path.join(self.label_dir, self.image_files[index].replace('.jpg', '.txt'))
    
    image = Image.open(img_path).convert("RGB")
    
    if os.path.exists(label_path):
        with open(label_path, 'r') as f:
            labels = f.readlines()
        labels = [line.strip().split() for line in labels]
        labels = [[int(label[0])] + list(map(float, label[1:])) for label in labels]
    else:
        labels = []
    
    if self.transform:
        image = self.transform(image)
    
    return image, labels
    
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    

    dataset = DisasterDetectionDataset(image_dir, label_dir, 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=batch_size, shuffle=True, num_workers=num_workers)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    
    return train_loader, val_loader
    
    
      
      
      
      
      
      
      
      
      
      
    

    def forward(self, x):
    x = self.backbone(x)
    x = self.head(x)
    return x
    
    
      
      
      
      
    

    model = YOLOv5(num_classes=2).to(device)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    train_loader, val_loader = get_data_loaders(image_dir, label_dir, batch_size=batch_size)
    
    writer = SummaryWriter()
    
    for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for images, labels in tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}"):
        images = images.to(device)
        labels = [torch.tensor(label, dtype=torch.float32).to(device) for label in labels]
    
        optimizer.zero_grad()
        outputs = model(images)
        loss = 0.0
        for i in range(len(labels)):
            if len(labels[i]) > 0:
                pred = outputs[i, :, :, :5 * len(labels[i])]
                pred = pred.permute(2, 0, 1).reshape(-1, 5)
                target = labels[i][:, 1:]
                loss += criterion(pred, target)
        loss.backward()
        optimizer.step()
    
        running_loss += loss.item()
    
    train_loss = running_loss / len(train_loader)
    writer.add_scalar('Training Loss', train_loss, epoch)
    
    model.eval()
    running_val_loss = 0.0
    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            labels = [torch.tensor(label, dtype=torch.float32).to(device) for label in labels]
    
            outputs = model(images)
            loss = 0.0
            for i in range(len(labels)):
                if len(labels[i]) > 0:
                    pred = outputs[i, :, :, :5 * len(labels[i])]
                    pred = pred.permute(2, 0, 1).reshape(-1, 5)
                    target = labels[i][:, 1:]
                    loss += criterion(pred, target)
    
            running_val_loss += loss.item()
    
    val_loss = running_val_loss / len(val_loader)
    writer.add_scalar('Validation Loss', val_loss, epoch)
    
    print(f"Epoch {epoch + 1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
    
    torch.save(model.state_dict(), "weights/best_model.pt")
    writer.close()
    
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    

    model = YOLOv5(num_classes=2).to(device)
    model.load_state_dict(torch.load(model_path))
    model.eval()
    
    _, val_loader = get_data_loaders(image_dir, label_dir)
    
    fig, axes = plt.subplots(num_samples, 2, figsize=(10, 5 * num_samples))
    with torch.no_grad():
    for i, (images, labels) in enumerate(val_loader):
        if i >= num_samples:
            break
    
        images = images.to(device)
        labels = [torch.tensor(label, dtype=torch.float32).to(device) for label in labels]
    
        outputs = model(images)
        outputs = outputs.cpu().numpy()
    
        images = images.cpu().numpy().transpose((0, 2, 3, 1))
        labels = [label.cpu().numpy() for label in labels]
    
        for j in range(len(images)):
            ax = axes[j] if num_samples > 1 else axes
            ax[0].imshow(images[j])
            ax[0].set_title(f"True: {labels[j]}")
            ax[0].axis('off')
    
            ax[1].imshow(images[j])
            ax[1].set_title(f"Predicted: {outputs[j]}")
            ax[1].axis('off')
    
    plt.tight_layout()
    plt.show()