first commit

voc2coco.py

@@ -0,0 +1,226 @@
+#!/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))