数据集裁剪VisDrone

visdrone无人机航拍数据集有10209张静态图像，标签格式为 .txt，类别为：

ignored regions(0)

pedestrian(1)           行人

people(2)                 人

bicycle(3)                 自行车

car(4)                       汽车

van(5)                      货车

truck(6)                    卡车

tricycle(7)                  三轮车

awning-tricycle(8)     遮阳三轮车

bus(9)                       公共汽车

motor(10)                 摩托车

others(11)

除去 0类别，有11类，整理的时候把0类别去掉（背景类）。

标签：

分别代表：<bbox_left>,<bbox_top>,<bbox_width>,<bbox_height>,

                  <score>,<object_category>,<截断率>,<遮挡率>

背景类(0) 的 score=0， 用这个可以过滤类别0

整理操作：

1. 过滤0类别，其他类别数减1，类别从0开始。
2. 坐标转换：绝对坐标转相对坐标
3. 选取其中一部分作为我们的数据集，太多了。（这里选了500张，train/val 是 4:1）

import os
import shutil
import random
from PIL import Image


# 新建数据集
data_path = "./VisDrone"
if not os.path.exists(data_path):
    os.mkdir(data_path)

an_path = os.path.join(data_path, "labels")
if not os.path.exists(an_path):
    os.mkdir(an_path)

im_path = os.path.join(data_path, "images")
if not os.path.exists(im_path):
    os.mkdir(im_path)

train_im_path = os.path.join(im_path, "train")
if not os.path.exists(train_im_path):
    os.mkdir(train_im_path)
val_im_path = os.path.join(im_path, "val")
if not os.path.exists(val_im_path):
    os.mkdir(val_im_path)
train_an_path = os.path.join(an_path, "train")
if not os.path.exists(train_an_path):
    os.mkdir(train_an_path)
val_an_path = os.path.join(an_path, "val")
if not os.path.exists(val_an_path):
    os.mkdir(val_an_path)

print("完成1")



raw_an_file = "/kaggle/input/visdrone/VisDrone/annotations"    # 原始数据集
raw_im_file = "/kaggle/input/visdrone/VisDrone/images"
dataset_size = 500                      # 选取数据集的尺寸
data_an_file = os.listdir(raw_an_file)[:dataset_size]

TRAIN_RATIO = 80        # 训练集/验证集 划分比例

def convert_box(size, box):
        # Convert VisDrone box to YOLO xywh box
        dw = 1. / size[0]
        dh = 1. / size[1]
        return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh


# 寻找对应raw_im_file[:500]
for i in range(0, len(data_an_file)):
    image_name = os.path.splitext(data_an_file[i])[0] + ".jpg"
    raw_image_path = os.path.join(raw_im_file, image_name) 
    if os.path.exists(raw_image_path):                # 存在对应图片, 另存标签和图片
        prob = random.randint(1, 100)                 # 按比例划分train/val dataset
        if prob < TRAIN_RATIO:                        # train dataset
            #shutil.copyfile(os.path.join(raw_an_file, data_an_file[i]),
            #               os.path.join(train_an_path, data_an_file[i]))
            shutil.copyfile(raw_image_path, 
                            os.path.join(train_im_path, image_name))
            # 标签格式转换
            img_size = Image.open(raw_image_path).size
            lines = []
            with open(os.path.join(raw_an_file, data_an_file[i]), "r") as file: # read annotation.txt
                for row in [x.split(',') for x in file.read().strip().splitlines()]:
                    if row[4] == '0':  # VisDrone 'ignored regions' class 0
                        continue
                    cls = int(row[5]) - 1  # 类别号-1
                    box = convert_box(img_size, tuple(map(int, row[:4])))
                    lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n")
                    
                    with open(os.path.join(train_an_path, data_an_file[i]), "w") as fl:
                        fl.writelines(lines)  # write label.txt
            
        else:    # val dataset
            #shutil.copyfile(os.path.join(raw_an_file, data_an_file[i]),
            #                os.path.join(val_an_path, data_an_file[i]))
            shutil.copyfile(raw_image_path, 
                            os.path.join(val_im_path, image_name))
            # 标签格式转换
            img_size = Image.open(raw_image_path).size
            lines = []
            with open(os.path.join(raw_an_file, data_an_file[i]), "r") as file: # read annotation.txt
                for row in [x.split(',') for x in file.read().strip().splitlines()]:
                    if row[4] == '0':  # VisDrone 'ignored regions' class 0
                        continue
                    cls = int(row[5]) - 1  # 类别号-1
                    box = convert_box(img_size, tuple(map(int, row[:4])))
                    lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}\n")
                    
                    with open(os.path.join(val_an_path, data_an_file[i]), "w") as fl:
                        fl.writelines(lines)  # write label.txt
        
    else:
        continue

注：

kaggle 文件下载：

打成压缩包：!tar -cf filename.tar folder