Convert yolo coordinates to VOC format

label.txt

Apple_porcelain
Beer
Book
Carbonated_drinks
Cup
Dispenser
Desk_lamp
Drugs
Glasses_case
Kettle
Milk
Mug
Metal_box
Nipper
Plate
Potato_chips
Seasoning_sauce
Spoon
Stick_snacks
Toy
Wash_product

part_label.txt

Beer_Cap
Carbonated_drinks_Cap
Cup_Cap
Cup_Mouth
Desk_lamp_Base
Desk_lamp_Lamp
Dispenser_Spout
Drugs_Cap
Kettle_Handle
Kettle_Lid
Kettle_Spout
Metal_box_Lid
Milk_Spout
Mug_Handle
Mug_Lid
Mug_Mouth
Nipper_Handle
Nipper_Head
Plate_Mouth
Potato_chips_Cap
Seasoning_sauce_Spout
Spoon_Scoop
Wash_product_Spout

voc_style_dir.py

import os
import glob
from sklearn.model_selection import train_test_split
from tqdm import tqdm

ROOT_PATH = '/data/picking'

xml_path = os.path.join(ROOT_PATH, 'train_annotation_voc')
annotations_path = os.path.join(ROOT_PATH, 'Annotations')
image_path = os.path.join(ROOT_PATH, 'JPEGImages')
txt_path = os.path.join(ROOT_PATH, 'ImageSets/Main')

def make_dir():
    try:
        os.mkdir('{}/JPEGImages'.format(ROOT_PATH))
    except:
        pass

    try:
        os.mkdir('{}/ImageSets'.format(ROOT_PATH))
    except:
        pass

    try:
        os.mkdir('{}/ImageSets/Main'.format(ROOT_PATH))
    except:
        pass

    try:
        os.mkdir('{}/Annotations'.format(ROOT_PATH))
    except:
        pass    

def move_file(files, tag):
    name_arr = []
    
    for t in tqdm(files):
        t_name = t.split('/')[-1]
        test_name = tag + '_' + t_name
        name_arr.append(test_name.split('.')[0])
        
        os.system('cp {} {}'.format(t, os.path.join(image_path, test_name)))
        # print('cp {} {}'.format(t, os.path.join(image_path, test_name)))

    return name_arr

def write_txt(names, tag):
    file_path = os.path.join(txt_path, tag)
    txt_name = file_path + '.txt'
    
    with open(txt_name, 'w') as f:
        for n in names:
            f.write(n + '\n')

def move_write(files, tag):
    print('move ', tag, ' images')
    names = move_file(files, tag)
    print('write ', tag, ' txt')
    write_txt(names, tag)
    
def move_annotation(files, tag):
    print('move ', tag, ' annotation')
    
    for item in tqdm(files):
        name = item.split('/')[-1].split('.')[0]
        xml_name = name + '.xml'
        annotation_name = tag + '_' + xml_name
        
        xml_file = os.path.join(xml_path, xml_name)
        annotation_file = os.path.join(annotations_path, annotation_name)
        
        os.system('cp {} {}'.format(xml_file, annotation_file))
    
def move_testset():
    test_files = glob.glob('{}/test/*.jp*'.format(ROOT_PATH))
    move_write(test_files, 'test')
    
def move_trainset():
    train_files = glob.glob('{}/train/*.jp*'.format(ROOT_PATH))
    
    train, val = train_test_split(train_files, test_size=0.2)
    
    move_write(train, 'train')
    move_write(val, 'val')
    
    move_annotation(train, 'train')
    move_annotation(val, 'val')

make_dir()    
move_trainset()
move_testset()

yolo_to_voc.py

# Script to convert yolo annotations to voc format

# Sample format
# <annotation>
#     <folder>_image_fashion</folder>
#     <filename>brooke-cagle-39574.jpg</filename>
#     <size>
#         <width>1200</width>
#         <height>800</height>
#         <depth>3</depth>
#     </size>
#     <segmented>0</segmented>
#     <object>
#         <name>head</name>
#         <pose>Unspecified</pose>
#         <truncated>0</truncated>
#         <difficult>0</difficult>
#         <bndbox>
#             <xmin>549</xmin>
#             <ymin>251</ymin>
#             <xmax>625</xmax>
#             <ymax>335</ymax>
#         </bndbox>
#     </object>
# <annotation>


import os
import glob
import xml.etree.cElementTree as ET
from PIL import Image

ANNOTATIONS_DIR_PREFIX = "3D_close_label/*"
IMAGES_DIR_PREFIX = "3D_close/"
DESTINATION_DIR = "3D_close_label_voc/"

CLASS_MAPPING = {
    '0': 'Apple_porcelain',
    '1': 'Beer',
    '2': 'Book', 
    '3': 'Carbonated_drinks', 
    '4': 'Cup', 
    '5': 'Dispenser', 
    '6': 'Desk_lamp', 
    '7': 'Drugs', 
    '8': 'Glasses_case',
    '9': 'Kettle',
    '10': 'Milk',
    '11': 'Mug',
    '12': 'Metal_box',
    '13': 'Nipper',
    '14': 'Plate',
    '15': 'Potato_chips',
    '16': 'Seasoning_sauce',
    '17': 'Spoon',
    '18': 'Stick_snacks',
    '19': 'Toy',
    '20': 'Wash_product'
}


def create_root(file_prefix, width, height):
    root = ET.Element("annotations")
    ET.SubElement(root, "filename").text = "{}.jpg".format(file_prefix)
    ET.SubElement(root, "folder").text = IMAGES_DIR_PREFIX
    size = ET.SubElement(root, "size")
    ET.SubElement(size, "width").text = str(width)
    ET.SubElement(size, "height").text = str(height)
    ET.SubElement(size, "depth").text = "3"
    return root


def create_object_annotation(root, voc_labels):
    for voc_label in voc_labels:
        obj = ET.SubElement(root, "object")
        ET.SubElement(obj, "name").text = voc_label[0]
        ET.SubElement(obj, "pose").text = "Unspecified"
        ET.SubElement(obj, "truncated").text = str(0)
        ET.SubElement(obj, "difficult").text = str(0)
        bbox = ET.SubElement(obj, "bndbox")
        ET.SubElement(bbox, "xmin").text = str(voc_label[1])
        ET.SubElement(bbox, "ymin").text = str(voc_label[2])
        ET.SubElement(bbox, "xmax").text = str(voc_label[3])
        ET.SubElement(bbox, "ymax").text = str(voc_label[4])
    return root


def create_file(file_prefix, width, height, voc_labels):
    root = create_root(file_prefix, width, height)
    root = create_object_annotation(root, voc_labels)
    tree = ET.ElementTree(root)
    tree.write("{}/{}.xml".format(DESTINATION_DIR, file_prefix))


def read_file(file_path):
    file_prefix = file_path.split("/")[-1].split('.')[0]
    image_file_name = "{}.jpg".format(file_prefix)
    img = Image.open("{}/{}".format(IMAGES_DIR_PREFIX, image_file_name))
    w, h = img.size
    
    with open(file_path, 'r') as file:
        lines = file.readlines()
        voc_labels = []
        
        for line in lines:
            voc = []
            line = line.strip()
            data = line.split()
            voc.append(CLASS_MAPPING.get(data[0]))
            bbox_width = float(data[3]) * w
            bbox_height = float(data[4]) * h
            center_x = float(data[1]) * w
            center_y = float(data[2]) * h
            voc.append(center_x - (bbox_width / 2))
            voc.append(center_y - (bbox_height / 2))
            voc.append(center_x + (bbox_width / 2))
            voc.append(center_y + (bbox_height / 2))
            voc_labels.append(voc)
        create_file(file_prefix, w, h, voc_labels)
        
    print("Processing complete for file: {}".format(file_path))

def start():
    if not os.path.exists(DESTINATION_DIR):
        os.makedirs(DESTINATION_DIR)
        
    for filename in glob.glob(ANNOTATIONS_DIR_PREFIX):
        if filename.endswith('txt'):
            read_file(filename)
        else:
            print("Skipping file: {}".format(filename))


if __name__ == "__main__":
    start()