226 lines
8.4 KiB
Python
Executable File
226 lines
8.4 KiB
Python
Executable File
#!/usr/bin/python
|
|
|
|
# pip install lxml
|
|
"""
|
|
需要先使用creat_coco_annotation.py生成 anno_train and anno_test 文件夹
|
|
"""
|
|
import sys
|
|
import os
|
|
import json
|
|
import xml.etree.ElementTree as ET
|
|
import glob
|
|
import shutil
|
|
import cv2
|
|
from tqdm import tqdm
|
|
from coco_visualization import checkCOCO
|
|
START_BOUNDING_BOX_ID = 1
|
|
PRE_DEFINE_CATEGORIES = None
|
|
# If necessary, pre-define category and its id
|
|
# PRE_DEFINE_CATEGORIES = {"aeroplane": 1, "bicycle": 2, "bird": 3, "boat": 4,
|
|
# "bottle":5, "bus": 6, "car": 7, "cat": 8, "chair": 9,
|
|
# "cow": 10, "diningtable": 11, "dog": 12, "horse": 13,
|
|
# "motorbike": 14, "person": 15, "pottedplant": 16,
|
|
# "sheep": 17, "sofa": 18, "train": 19, "tvmonitor": 20}
|
|
|
|
|
|
def get(root, name):
|
|
vars = root.findall(name)
|
|
return vars
|
|
|
|
|
|
def get_and_check(root, name, length):
|
|
vars = root.findall(name)
|
|
if len(vars) == 0:
|
|
raise ValueError("Can not find %s in %s." % (name, root.tag))
|
|
if length > 0 and len(vars) != length:
|
|
raise ValueError(
|
|
"The size of %s is supposed to be %d, but is %d."
|
|
% (name, length, len(vars))
|
|
)
|
|
if length == 1:
|
|
vars = vars[0]
|
|
return vars
|
|
|
|
|
|
def get_filename_as_int(filename):
|
|
try:
|
|
filename = filename.replace("\\", "/")
|
|
filename = os.path.splitext(os.path.basename(filename))[0]
|
|
return str(filename)
|
|
except:
|
|
raise ValueError("Filename %s is supposed to be an integer." % (filename))
|
|
|
|
|
|
def get_categories(xml_files):
|
|
"""Generate category name to id mapping from a list of xml files.
|
|
|
|
Arguments:
|
|
xml_files {list} -- A list of xml file paths.
|
|
|
|
Returns:
|
|
dict -- category name to id mapping.
|
|
"""
|
|
classes_names = []
|
|
for xml_file in xml_files:
|
|
tree = ET.parse(xml_file)
|
|
root = tree.getroot()
|
|
for member in root.findall("object"):
|
|
classes_names.append(member[0].text)
|
|
classes_names = list(set(classes_names))
|
|
classes_names.sort()
|
|
#return {name: i for i, name in enumerate(classes_names)}
|
|
return {"bird": 1}
|
|
|
|
|
|
def convert(rootpath,xml_files, json_file):
|
|
json_dict = {"images": [], "type": "instances", "annotations": [], "categories": []}
|
|
if PRE_DEFINE_CATEGORIES is not None:
|
|
categories = PRE_DEFINE_CATEGORIES
|
|
else:
|
|
categories = get_categories(xml_files)
|
|
bnd_id = START_BOUNDING_BOX_ID
|
|
for xml_file in xml_files:
|
|
tree = ET.parse(xml_file)
|
|
root = tree.getroot()
|
|
path = get(root, "path")
|
|
if len(path) == 1:
|
|
filename = os.path.basename(xml_file).replace(".xml","")
|
|
elif len(path) == 0:
|
|
filename = os.path.basename(xml_file).replace(".xml", "")
|
|
else:
|
|
raise ValueError("%d paths found in %s" % (len(path), xml_file))
|
|
## The filename must be a number
|
|
image_id = get_filename_as_int(filename)
|
|
try:
|
|
size = get_and_check(root, "size", 1)
|
|
width = int(get_and_check(size, "width", 1).text)
|
|
height = int(get_and_check(size, "height", 1).text)
|
|
except:
|
|
imgPath = os.path.join(rootpath,"JPEGImages","{}.jpg".format(str(filename)))
|
|
try:
|
|
img = cv2.imread(imgPath)
|
|
|
|
except:
|
|
imgPath = os.path.join(rootpath,"JPEGImages","{}.png".format(str(filename)))
|
|
img = cv2.imread(imgPath)
|
|
|
|
width =int(img.shape[1])
|
|
height = int(img.shape[0])
|
|
if not os.path.exists(os.path.join(rootpath,"JPEGImages","{}.jpg".format(str(filename)))): #图片里面既有jpg又有png
|
|
image = {
|
|
"file_name": "{}.png".format(filename),
|
|
"height": height,
|
|
"width": width,
|
|
"id": image_id,
|
|
}
|
|
else:
|
|
image = {
|
|
"file_name": "{}.jpg".format(filename),
|
|
"height": height,
|
|
"width": width,
|
|
"id": image_id,
|
|
}
|
|
json_dict["images"].append(image)
|
|
## Currently we do not support segmentation.
|
|
# segmented = get_and_check(root, 'segmented', 1).text
|
|
# assert segmented == '0'
|
|
#print(filename)
|
|
for obj in get(root, "object"):
|
|
category = "bird"#get_and_check(obj, "name", 1).text
|
|
# if category not in categories:
|
|
# new_id = len(categories)
|
|
# categories[category] = new_id
|
|
category_id = categories[category]
|
|
bndbox = get_and_check(obj, "bndbox", 1)
|
|
xmin = int(get_and_check(bndbox, "xmin", 1).text)
|
|
ymin = int(get_and_check(bndbox, "ymin", 1).text)
|
|
xmax = int(get_and_check(bndbox, "xmax", 1).text)
|
|
ymax = int(get_and_check(bndbox, "ymax", 1).text)
|
|
assert xmax > xmin
|
|
assert ymax > ymin
|
|
o_width = abs(xmax - xmin)
|
|
o_height = abs(ymax - ymin)
|
|
ann = {
|
|
"area": o_width * o_height,
|
|
"iscrowd": 0,
|
|
"image_id": image_id,
|
|
"bbox": [xmin, ymin, o_width, o_height],
|
|
"category_id": category_id,
|
|
"id": bnd_id,
|
|
"ignore": 0,
|
|
"segmentation": [],
|
|
}
|
|
json_dict["annotations"].append(ann)
|
|
bnd_id = bnd_id + 1
|
|
|
|
for cate, cid in categories.items():
|
|
cat = {"supercategory": "none", "id": cid, "name": cate}
|
|
json_dict["categories"].append(cat)
|
|
|
|
os.makedirs(os.path.dirname(json_file), exist_ok=True)
|
|
json_fp = open(json_file, "w")
|
|
json_str = json.dumps(json_dict)
|
|
json_fp.write(json_str)
|
|
json_fp.close()
|
|
def data_pre(root):
|
|
xml_root = os.path.join(root,"ImageSets/Main/train.txt")
|
|
result_root = os.path.join(root,"COCO/Anno_train")
|
|
if not os.path.exists(result_root):
|
|
os.makedirs(result_root)
|
|
with open(xml_root) as f:
|
|
lines = f.readlines()
|
|
for line in tqdm(lines):
|
|
sourcePath = os.path.join(root,"Annotations","{}.xml".format(str(line).replace("\n","")))
|
|
desPath = os.path.join(result_root,"{}.xml".format(str(line)).replace("\n",""))
|
|
shutil.copy(sourcePath,desPath)
|
|
print("train set is ok")
|
|
xml_root = os.path.join(root,"ImageSets/Main/test.txt")
|
|
result_root = os.path.join(root,"COCO/Anno_val")
|
|
if not os.path.exists(result_root):
|
|
os.makedirs(result_root)
|
|
with open(xml_root) as f:
|
|
lines = f.readlines()
|
|
for line in tqdm(lines):
|
|
sourcePath = os.path.join(root,"Annotations","{}.xml".format(str(line).replace("\n","")))
|
|
desPath = os.path.join(result_root,"{}.xml".format(str(line)).replace("\n",""))
|
|
shutil.copy(sourcePath,desPath)
|
|
print("val set is ok")
|
|
if __name__ == "__main__":
|
|
#os.system("sudo chmod 777 -R /data/DataManagement/Project/CAS_ProjectBird/CDB_ProjectBird/")
|
|
#root = "/data/DataManagement/Project/CAS_ProjectBird/CDB_ProjectBird/"
|
|
#data_pre(root)
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Convert Pascal VOC annotation to COCO format."
|
|
)
|
|
parser.add_argument("root", default="/data4/Hatcher/smallobject/voc_class/CDB_airplane/",help="root", type=str)
|
|
|
|
args = parser.parse_args()
|
|
root = args.root
|
|
jsonPath = os.path.join(root,"COCO/annotations")
|
|
if not os.path.exists(jsonPath):
|
|
os.makedirs(jsonPath)
|
|
jsonTrain = os.path.join(jsonPath,"train.json")
|
|
jsonval = os.path.join(jsonPath, "val.json")
|
|
result_root = os.path.join(root, "COCO/Anno_train")
|
|
|
|
xml_files = glob.glob(os.path.join(result_root, "*.xml"))
|
|
# If you want to do train/val split, you can pass a subset of xml files to convert function.
|
|
print("Number of train xml files: {}".format(len(xml_files)))
|
|
convert(root,xml_files, jsonTrain)
|
|
print("train COCO is ok")
|
|
|
|
result_root = os.path.join(root, "COCO/Anno_test")
|
|
xml_files = glob.glob(os.path.join(result_root, "*.xml"))
|
|
# If you want to do train/val split, you can pass a subset of xml files to convert function.
|
|
print("Number of val xml files: {}".format(len(xml_files)))
|
|
convert(root,xml_files, jsonval)
|
|
print("val COCO is ok")
|
|
if not os.path.join(root,"COCO/train"):
|
|
os.makedirs(os.path.join(root,"COCO/train"))
|
|
if not os.path.join(root,"COCO/val"):
|
|
os.makedirs(os.path.join(root,"COCO/val"))
|
|
os.system("ln -s {}/images/train {}/COCO/train".format(root,root))
|
|
os.system("ln -s {}/images/test {}/COCO/val".format(root, root))
|
|
checkCOCO(jsonTrain,"{}/COCO/train".format(root)) |