MapillaryVistas数据集入门

pythonCopy codeimport json
# 读取标注文件
annotation_file = "annotations.json"
with open(annotation_file, "r") as f:
    annotations = json.load(f)
# 遍历每个标注
for annotation in annotations:
    image_path = annotation["image_path"]
    objects = annotation["objects"]
    
    # 在图像上绘制边框和标签
    for obj in objects:
        category = obj["category"]
        bbox = obj["bbox"]
        
        # 在图像上绘制矩形框
        draw_rectangle(image_path, bbox)
        
        # 在图像上绘制标签
        draw_text(image_path, category)

pythonCopy codeimport torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
# 自定义数据集类
class MapillaryVistasDataset(Dataset):
    def __init__(self, image_paths, label_paths, transform=None):
        self.image_paths = image_paths
        self.label_paths = label_paths
        self.transform = transform
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        # 读取图像和标签
        image = Image.open(self.image_paths[idx]).convert("RGB")
        label = Image.open(self.label_paths[idx]).convert("L")
        
        if self.transform:
            image = self.transform(image)
            label = self.transform(label)
        
        return image, label
# 定义自定义转换
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])
# 定义训练集和验证集路径
train_image_paths = ["train/1.jpg", "train/2.jpg", ...]
train_label_paths = ["train/1_label.png", "train/2_label.png", ...]
val_image_paths = ["val/1.jpg", "val/2.jpg", ...]
val_label_paths = ["val/1_label.png", "val/2_label.png", ...]
# 创建训练集和验证集实例
train_dataset = MapillaryVistasDataset(train_image_paths, train_label_paths, transform)
val_dataset = MapillaryVistasDataset(val_image_paths, val_label_paths, transform)
# 创建数据加载器
train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=16, shuffle=False)
# 定义模型
model = UNet()
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 开始训练过程
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    train_loss = 0.0
    
    for images, labels in train_dataloader:
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item()
    
    # 打印训练损失
    print(f"Epoch {epoch+1}/{num_epochs} - Train Loss: {train_loss/len(train_dataloader)}")
    
    # 在验证集上评估模型
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for images, labels in val_dataloader:
            outputs = model(images)
            loss = criterion(outputs, labels)
            val_loss += loss.item()
    # 打印验证损失
    print(f"Epoch {epoch+1}/{num_epochs} - Validation Loss: {val_loss/len(val_dataloader)}")
# 模型训练完成后，可以使用该模型对新的图像进行语义分割预测
test_image = Image.open("test.jpg").convert("RGB")
test_image = transform(test_image).unsqueeze(0)
prediction = model(test_image)
segmentation_mask = torch.argmax(prediction, dim=1)
# 显示预测结果
plt.imshow(segmentation_mask.squeeze().detach().numpy(), cmap='jet')
plt.show()