coco格式伪标签生成代码

coco格式伪标签生成代码

注：name_class 要更换为自己数据集的标签

import json
import os
from glob import glob
from tqdm import tqdm
from PIL import Image
name_class =['liner', 'bulk carrier', 'warship', 'sailboat']
label_ids = {name: i  for i, name in enumerate(name_class)}


def save(images, annotations):
    ann = {}
    ann['type'] = 'instances'
    ann['images'] = images
    ann['annotations'] = annotations

    categories = []
    for k, v in label_ids.items():
        categories.append({"name": k, "id": v})
    ann['categories'] = categories
    json.dump(ann, open('images/instances_test.json', 'w'))


def test_dataset(im_dir):
    # im_list = glob(os.path.join(im_dir, '*.jpg'))
    im_list = os.listdir(im_dir)

    im_list.sort(key=lambda x: int(x.split('.')[0]))
    idx = 0
    image_id = -1
    images = []
    annotations = []
    for im_path in tqdm(im_list):
        image_id += 1
        im = Image.open(os.path.join(im_dir,im_path))
        w, h = im.size
        image = {'file_name': os.path.basename(im_path), 'width': w, 'height': h, 'id': image_id}
        images.append(image)
        labels = [[10, 10, 20, 20]]
        for label in labels:
            bbox = [label[0], label[1], label[2] - label[0], label[3] - label[1]]
            seg = []
            ann = {'segmentation': [seg], 'area': bbox[2] * bbox[3], 'iscrowd': 0, 'image_id': image_id,
                   'bbox': bbox, 'category_id': 1, 'id': idx, 'ignore': 0}
            idx += 1
            annotations.append(ann)
    save(images, annotations)


if __name__ == '__main__':
    test_dir = 'images/'
    print("generate test json label file.")
    test_dataset(test_dir)