shuoshuoshuo
/
labelme2yolo
9
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

use multiprocessing to boost converting speed

This commit is contained in:
Wang Xin 2022-07-27 12:19:14 +08:00
parent 3cb919641a
commit fe3af9b104
1 changed files with 16 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
labelme2yolo.py
View File

@ -10,6 +10,7 @@ import argparse
import shutil import shutil
import math import math
from collections import OrderedDict from collections import OrderedDict
from multiprocessing import Pool
import json import json
import cv2 import cv2
@ -101,26 +102,32 @@ class Labelme2YOLO(object):
# also get image from labelme json file and save them under images folder # also get image from labelme json file and save them under images folder
for target_dir, json_names in zip(('train/', 'val/', 'test/'), for target_dir, json_names in zip(('train/', 'val/', 'test/'),
(train_json_names, val_json_names, test_json_names)): (train_json_names, val_json_names, test_json_names)):
pool = Pool(os.cpu_count() - 1)
for json_name in json_names: for json_name in json_names:
json_path = os.path.join(self._json_dir, json_name) pool.apply_async(self.covert_json_to_text, args=(target_dir, json_name))
json_data = json.load(open(json_path)) pool.close()
pool.join()
print('Converting %s for %s ...' % (json_name, target_dir.replace('/', ''))) print('Generating dataset.yaml file ...')
self._save_dataset_yaml()
img_path = self._save_yolo_image(json_data, def covert_json_to_text(self, target_dir, json_name):
json_path = os.path.join(self._json_dir, json_name)
json_data = json.load(open(json_path))
print('Converting %s for %s ...' % (json_name, target_dir.replace('/', '')))
img_path = self._save_yolo_image(json_data,
json_name, json_name,
self._image_dir_path, self._image_dir_path,
target_dir) target_dir)
yolo_obj_list = self._get_yolo_object_list(json_data, img_path) yolo_obj_list = self._get_yolo_object_list(json_data, img_path)
self._save_yolo_label(json_name, self._save_yolo_label(json_name,
self._label_dir_path, self._label_dir_path,
target_dir, target_dir,
yolo_obj_list) yolo_obj_list)
print('Generating dataset.yaml file ...')
self._save_dataset_yaml()
def convert_one(self, json_name): def convert_one(self, json_name):
json_path = os.path.join(self._json_dir, json_name) json_path = os.path.join(self._json_dir, json_name)
json_data = json.load(open(json_path)) json_data = json.load(open(json_path))