#!/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))