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()