eriklindernoren / PyTorch-YOLOv3

Minimal PyTorch implementation of YOLOv3
GNU General Public License v3.0
7.32k stars 2.63k forks source link

VOC #519

Open Muxindawang opened 4 years ago

Muxindawang commented 4 years ago

I want to ask how to train e on VOC dataset? Whether I need to write my own configuration file?

MinaGabriel commented 4 years ago

I think you should convert VOC to YOLO on your own I used the following code to rebuild my everything customer configuration all you need is the images and the annotation path (which is VOC)


import os
from datetime import datetime
from shutil import copy
import xml.etree.ElementTree as ET
from absl import app, flags, logging

FLAGS = flags.FLAGS
flags.DEFINE_string('images', None, 'train images folder')
flags.DEFINE_string('annotation', None, 'train annotation folder')
flags.DEFINE_bool('rebuild', True, 'true if not all images has equivalent annotation files')
flags.DEFINE_bool('VOCtoYOLO', False, 'add a new directory with Yolo annotation from VOC')
flags.DEFINE_string('classes', None, 'classes file location, .names file in root directory')
flags.DEFINE_bool('new_name', True, 'Change images and annotation file name')

flags.mark_flag_as_required('images')  # path to train images………
flags.mark_flag_as_required('annotation')  # path to train annotation

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def main(argv):
    images = FLAGS.images
    annotation = FLAGS.annotation
    rebuild = FLAGS.rebuild
    VOCtoYOLO = FLAGS.VOCtoYOLO

    # create new directory for validation anyway
    now = datetime.now()
    # create two new folders
    new_dir = 'images-' + now.strftime('%m-%d-%Y-%H-%M-%S')
    os.makedirs(new_dir)

    classes = []
    if FLAGS.classes is not None:
        # copy the classes file it self to the new directory
        copy(FLAGS.classes,
             new_dir)
        # create the classes list
        with open(FLAGS.classes) as f:
            classes = [my_class.strip() for my_class in f]

    if rebuild:

        new_dir_images = new_dir + '/images'
        new_dir_annotation = new_dir + '/annotation'
        yolo_dir = new_dir + '/YOLO_annotation'

        # create two folders in the new images: images and annotation
        os.makedirs(new_dir_images)
        os.makedirs(new_dir_annotation)
        if VOCtoYOLO:
            os.makedirs(yolo_dir)

        counter = 0
        if os.path.exists(annotation):

            # count the number of annotation files
            annotation_count = len([file for file in os.listdir(annotation)])
            train_count = annotation_count - (annotation_count // 3)

            # create train and validation files
            train = open(os.path.join(new_dir, 'train.txt'), 'w')
            validation = open(os.path.join(new_dir, 'validation.txt'), 'w')

            for annotation_file_name in os.listdir(annotation):
                annotation_file = os.path.splitext(annotation_file_name)[0]
                image_file_name = annotation_file + '.JPG'

                # if image exists move both annotation and images to new directory
                if os.path.exists(os.path.join(images, image_file_name)):
                    # copy annotation
                    copy(os.path.join(annotation, annotation_file_name),
                         os.path.join(new_dir_annotation, annotation_file_name))
                    copy(os.path.join(images, image_file_name),
                         os.path.join(new_dir_images, image_file_name))

                    # print('copying ' + annotation_file_name + ' & ' + image_file_name + ' >>> ' + new_dir)
                    counter += 1

                    if counter < train_count:
                        train.write('data/custom/images/' + image_file_name + '\n')
                    else:
                        validation.write('data/custom/images/' + image_file_name + '\n')

                    # if convert VOC to YOLO
                    if VOCtoYOLO and classes:
                        # create the new Yolo file
                        yolo_file = open(os.path.join(yolo_dir, annotation_file) + '.txt', 'w')
                        tree = ET.parse(os.path.join(annotation, annotation_file_name))
                        root = tree.getroot()
                        size = root.find('size')
                        w = int(size.find('width').text)
                        h = int(size.find('height').text)
                        for obj in root.iter('object'):
                            cls = obj.find('name').text
                            cls_id = classes.index(cls)
                            xmlbox = obj.find('bndbox')
                            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
                                 float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
                            bb = convert((w, h), b)

                            yolo_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

        logging.info('Copied ' + str(counter) + ' Images and Annotation')

if __name__ == '__main__':
    app.run(main)
Muxindawang commented 4 years ago

Do you mean this code converts VOC to COCO? then I can train directly (I just need images with labels)

MinaGabriel commented 4 years ago

should convert VOC to YOLO like annotation <class id> <x> <y> <w> <h>

Muxindawang commented 4 years ago

xml convert to json? Txt needs to be made by myself. right?

risemeup commented 4 years ago

should convert VOC to YOLO like annotation <class id> <x> <y> <w> <h>

hello,have you train voc sussesflu? how about map?

Muxindawang commented 4 years ago

should convert VOC to YOLO like annotation <class id> <x> <y> <w> <h>

hello,have you train voc sussesflu? how about map?

I found a label tool which can make yolo format, so I give up voc

samcw commented 4 years ago

should convert VOC to YOLO like annotation <class id> <x> <y> <w> <h>

hello,have you train voc sussesflu? how about map?

I found a label tool which can make yolo format, so I give up voc

Hello, can you share this tool, I have the same trouble like yours. Thanks.

risemeup commented 4 years ago

sorry,I don't use this code finally.It's not difficult to make new annotation from xml format by yourself.Yout can try it!