Created
October 2, 2025 07:22
-
-
Save sugiana/b30d4e48c84d71f8cf47980b718f7813 to your computer and use it in GitHub Desktop.
Buat YAML untuk YOLO
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # Kebutuhan: env/bin/pip install ruamel.yaml | |
| # Ini untuk mengekstrak file ZIP hasil export berformat YOLO dari annotation | |
| # tool bernama Label Studio. Nanti akan terbentuk file YAML sebagai file | |
| # konfigurasi pelatihan: | |
| # env/bin/yolo train data=latihan.yaml | |
| import sys | |
| import os | |
| import shutil | |
| from zipfile import ZipFile | |
| from random import choice | |
| from argparse import ArgumentParser | |
| from ruamel.yaml import YAML | |
| def mkdir(name): | |
| print(f'mkdir {name}') | |
| os.mkdir(name) | |
| def move(file, path_dir): | |
| print(f'mv {file} {path_dir}') | |
| shutil.move(file, path_dir) | |
| def files(image_file): | |
| name, _ = os.path.splitext(image_file) | |
| label_file = name + '.txt' | |
| label_file = os.path.join(labels_dir, label_file) | |
| image_file = os.path.join(images_dir, image_file) | |
| return image_file, label_file | |
| yaml_file = 'latihan.yaml' | |
| help_yaml = f'default {yaml_file}' | |
| pars = ArgumentParser() | |
| pars.add_argument('zipfile') | |
| pars.add_argument('--yaml-file', default=yaml_file, help=help_yaml) | |
| option = pars.parse_args(sys.argv[1:]) | |
| root_dir, _ = os.path.splitext(option.zipfile) | |
| with ZipFile(option.zipfile) as z: | |
| z.extractall(root_dir) | |
| classes_file = os.path.join(root_dir, 'classes.txt') | |
| no = -1 | |
| names = dict() | |
| with open(classes_file) as f: | |
| for line in f.readlines(): | |
| no += 1 | |
| names[no] = line.rstrip() | |
| yaml = YAML() | |
| data = dict( | |
| names=names, | |
| path=os.path.realpath(root_dir), | |
| train='images/train', | |
| val='images/valid') | |
| print('Buat file', option.yaml_file) | |
| with open(option.yaml_file, 'w') as f: | |
| yaml.dump(data, f) | |
| images_dir = os.path.join(root_dir, 'images') | |
| labels_dir = os.path.join(root_dir, 'labels') | |
| images_train_dir = os.path.join(images_dir, 'train') | |
| labels_train_dir = os.path.join(labels_dir, 'train') | |
| images_valid_dir = os.path.join(images_dir, 'valid') | |
| labels_valid_dir = os.path.join(labels_dir, 'valid') | |
| if os.path.exists(images_train_dir): | |
| print(f'Direktori {images_train_dir} sudah ada. Akhiri.') | |
| sys.exit(1) | |
| mkdir(images_train_dir) | |
| mkdir(labels_train_dir) | |
| mkdir(images_valid_dir) | |
| mkdir(labels_valid_dir) | |
| image_files = os.listdir(images_dir) | |
| count = len(image_files) | |
| # 20% file untuk validation selama proses pelatihan | |
| valid_count = int(count/5) | |
| for _ in range(valid_count): | |
| image_file = choice(image_files) | |
| image_file, label_file = files(image_file) | |
| if os.path.isdir(image_file): | |
| continue | |
| move(image_file, images_valid_dir) | |
| move(label_file, labels_valid_dir) | |
| for image_file in os.listdir(images_dir): | |
| image_file, label_file = files(image_file) | |
| if os.path.isdir(image_file): | |
| continue | |
| move(image_file, images_train_dir) | |
| move(label_file, labels_train_dir) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment