coco数据集提取People 目标检测数据构建darknet 数据集

from pycocotools.coco import COCO
import numpy as np
import skimage.io as io
import matplotlib.pyplot as plt
import os
import cv2
import shutil

# 要绝对路径
dataDir='/home/haifan/haifan/zhouyong/datasets/coco/annotations_train2017'
dataType='train2017'
annFile='{}/annotations/instances_{}.json'.format(dataDir,dataType)

imgs_dir = "/home/haifan/haifan/zhouyong/datasets/coco/train2017/"

# 初始化标注数据的 COCO api
coco=COCO(annFile)


#获取person 相关的所有IDS
catIds = coco.getCatIds(catNms=['person'])
print("catIds = ", catIds)
imgIds = coco.getImgIds(catIds=catIds)

imgs = coco.loadImgs(imgIds)

output_train_file = "people/train_people.txt"
output_test_file = "people/test_people.txt"
label_dir = "people/txt/"
image_dir = "people/"

f_train = open(output_train_file, "w")
f_test = open(output_test_file, "w")
count = 0
for img in imgs:
    if count < 10000:
        fl = open(label_dir + img["file_name"][:-4] + ".txt", "w")
        line = ""
        annIds = coco.getAnnIds(imgIds=imgIds[count])
        # print("annIds = ", annIds)
        anns = coco.loadAnns(annIds)
        print("count = ", count , img["file_name"])
        # img_dat = cv2.imread(imgs_dir + img["file_name"])
        height, width = img["height"], img["width"]

        # 9:1 数据用来训练，剩余数据用来测试
        if count < 9000:
            f_train.write("data/people/images/" + img["file_name"] + "\n")
        else:
            f_test.write("data/people/images/" + img["file_name"] + "\n")

        for ann in anns:
            # print(ann)
            if ann['category_id'] == 1:
                x, y, w, h = tuple(ann["bbox"])
                x_center = x + w/2
                y_center = y + h/2
                x_center = 1.0 * x_center / width
                y_center = 1.0 * y_center / height
                w = 1.0 * w / width
                h = 1.0 * h / height
                line += str(0) + " " + str(x_center) + " " + str(y_center) + " " + str(w) + " " + str(h) + "\n"
                # bbox [x, y, w, h]
                # cv2.rectangle(img_dat, (int(x), int(y)), (int(x + w), int(y + h)), (255, 0, 0), 2)
        # cv2.imwrite("test_" + str(count) + ".jpg", img_dat)
        shutil.copy(imgs_dir + img["file_name"], image_dir)
        fl.write(line)
        fl.close()
    else:
        break
    count += 1

f_train.close()
f_test.close()