Convert Dataset to COCO Format

convert_to_coco.py

import json
import os

TRAIN_PATH = 'PCBData/PCBData/trainval.txt'
TEST_PATH = 'PCBData/PCBData/test.txt'


def create_data(data_path, output_path):
    images = []
    anns = []
    
    with open(data_path, 'r') as f:
        data = f.read().splitlines()
    dataset = []
    counter = 0
    for idx, example in enumerate(data):
        image_path, annotations_path = example.split()
        image_path = os.path.join('PCBData', 'PCBData', image_path.replace('.jpg', '_test.jpg'))
        annotations_path = os.path.join('PCBData', 'PCBData', annotations_path)

        with open(annotations_path, 'r') as f:
            annotations = f.read().splitlines()
            
        for ann in annotations:
            x, y, x2, y2 = ann.split()[:-1]
            
            anns.append({
                'image_id': idx,
                'iscrowd': 0,
                'area': (int(x2)-int(x)) * (int(y2)-int(y)),
                'category_id': int(ann.split()[-1])-1,
                'bbox': [int(x), int(y), int(x2)-int(x), int(y2)-int(y)],
                'id': counter
            })

            counter += 1

        images.append({
            'file_name': image_path,
            'width': 640,
            'height': 640,
            'id': idx
        })
        
    dataset = {
        'images': images,
        'annotations': anns,
        'categories': [
            {'id': 0, 'name': 'open'},
            {'id': 1, 'name': 'short'},
            {'id': 2, 'name': 'mousebite'},
            {'id': 3, 'name': 'spur'},
            {'id': 4, 'name': 'copper'},
            {'id': 5, 'name': 'pin-hole'},
        ]
    }
    with open(output_path, 'w') as f:
        json.dump(dataset, f)

create_data(TRAIN_PATH, 'train.json')
create_data(TEST_PATH, 'test.json')