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

improve convertion speed (#44)

This commit is contained in:
Wang Xin 2024-02-29 20:28:53 +08:00 committed by GitHub
parent f777359f33
commit 1e697c78e9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/labelme2yolo/l2y.py
View File

@ -14,6 +14,7 @@ import random
import shutil import shutil
import uuid import uuid
import logging import logging
from functools import partial
from multiprocessing import Pool from multiprocessing import Pool
import PIL.ExifTags import PIL.ExifTags
@ -21,7 +22,7 @@ import PIL.Image
import PIL.ImageOps import PIL.ImageOps
import cv2 import cv2
import numpy as np import numpy as np
from rich.progress import track from rich.progress import Progress
# set seed # set seed
random.seed(12345678) random.seed(12345678)
@ -245,9 +246,11 @@ class Labelme2YOLO:
names = (train_json_names, val_json_names, test_json_names) names = (train_json_names, val_json_names, test_json_names)
for target_dir, json_names in zip(dirs, names): for target_dir, json_names in zip(dirs, names):
logger.info("Converting %s set ...", target_dir) logger.info("Converting %s set ...", target_dir)
with Pool(os.cpu_count() - 1) as pool: with Pool(os.cpu_count() - 1) as pool, Progress() as progress:
for json_name in track(json_names): task = progress.add_task("[cyan]Converting...", total=len(json_names))
pool.apply(self.covert_json_to_text, (target_dir, json_name)) func = partial(self.covert_json_to_text, target_dir)
for _ in pool.map(func, json_names):
progress.update(task, advance=1)
self._save_dataset_yaml() self._save_dataset_yaml()