如何yolov8训练使用——电塔电线电线杆缺陷检测数据集 10000张 带标注 voc yolo 电线杆子缺陷数据集输电线塔缺陷数据集

    for xml_file in os.listdir(voc_dir):
    if not xml_file.endswith('.xml'):
        continue
    
    tree = ET.parse(os.path.join(voc_dir, xml_file))
    root = tree.getroot()
    
    image_width = int(root.find('size/width').text)
    image_height = int(root.find('size/height').text)
    
    yolo_file = os.path.join(yolo_dir, xml_file.replace('.xml', '.txt'))
    with open(yolo_file, 'w') as f:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            if class_name not in class_names:
                continue
            class_id = class_names.index(class_name)
    
            bndbox = obj.find('bndbox')
            xmin = int(bndbox.find('xmin').text)
            ymin = int(bndbox.find('ymin').text)
            xmax = int(bndbox.find('xmax').text)
            ymax = int(bndbox.find('ymax').text)
    
            x_center = (xmin + xmax) / 2.0 / image_width
            y_center = (ymin + ymax) / 2.0 / image_height
            width = (xmax - xmin) / image_width
            height = (ymax - ymin) / image_height
    
            f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")
    
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    

    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 = PowerLineDefectDetectionDataset(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
    
    
      
      
      
      
      
      
      
      
      
      
    

    model = YOLO('yolov8n.pt')  # Load a pretrained YOLOv8 model
    model = model.to(device)
    
    train_loader, val_loader = get_data_loaders(image_dir, label_dir, batch_size=batch_size)
    
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    criterion = torch.nn.MSELoss()
    
    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}")
    
    model.save("weights/best_model.pt")
    writer.close()
    
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    

    model = YOLO(model_path)
    model = model.to(device)
    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()