darknet标签转化为COCO标签

import sys
import json
import cv2
import os
import shutil

dataset = { "info": {
            "description": "XXX in COCO dataset.", 
            "url": "", 
            "version": "1.0", 
            "year": , 
            "contributor": "aimhabo", 
            "date_created": "2019-03-25"}, 
            "images":[],
            "annotations":[],
            "categories": [
            {"supercategory:": "person", "id": 1, "name": "person"},
            {"supercategory:": "car", "id": 2, "name": "car"},
            {"supercategory:": "truck", "id": 3, "name": "truck"},
            {"supercategory:": "bus", "id": 4, "name": "bus"}
            ]
}

datapath = "scripts/images"
annopath = "scripts/labels"
trainsetfile = "scripts/trainimage.list"
outputpath = "scripts"
phase = "XXXTrainCOCO"
classes = {"background": 0, "person": 1, "car": 2, "truck": 3, "bus": 4}

with open(trainsetfile) as f:
    count = 1
    cnt = 0
    annoid = 0
    for line in f:
        cnt += 1
        line = line.strip()

        name,ext = os.path.basename(line).split('.')
        
        imagepath = os.path.join(datapath, name + ".jpg")
        # no obstacle currently drop it
        txtpath = os.path.join(annopath, name + ".txt")
        if not os.path.exists(txtpath):
            print txtpath
            continue

        im = cv2.imread(imagepath)
            
        height, width, channels = im.shape

        if cnt % 1000 == 0:
            print cnt
                
        dataset["images"].append({"license": 5, "file_name": line, "coco_url": "local", "height": height, "width": width, "flickr_url": "local", "id": cnt})
        with open(txtpath) as annof:
            annos = annof.readlines()
        
        for ii, anno in enumerate(annos):
            parts = anno.strip().split(' ')
            if len(parts) is not 5:
                continue
            class_id = int(parts[0])
            x = float(parts[1])
            y = float(parts[2])
            w = float(parts[3])
            h = float(parts[4])
            if parts[0].find("group") == -1:
                iscrowd = 0
            else:
                iscrowd = 1

            annoid = annoid + 1

            class_id += 1 # start from 1 instead of 0
            x1 = int((x-w/2)*width)
            y1 = int((y-h/2)*height)
            wid = int(w*width)
            hei = int(h*height)
            
            dataset["annotations"].append({
                "segmentation": [],
                "iscrowd": iscrowd,
                "area": wid * hei,
                "image_id": cnt,
                "bbox": [x1, y1, wid, hei],
                "category_id": class_id,
                "id": annoid
            })
        count += 1

json_name = os.path.join(outputpath, "{}.json".format(phase))

with open(json_name, 'w') as f:
    json.dump(dataset, f)

 

posted @ 2019-06-04 16:17  aimhabo  阅读(313)  评论(0编辑  收藏