Spaces:
Runtime error
Runtime error
File size: 2,399 Bytes
4f8ad24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import os.path
from typing import Iterator
from hbutils.system import remove
from tqdm.auto import tqdm
from ..model import ImageItem
from ..utils import get_task_names
class BaseExporter:
def pre_export(self):
raise NotImplementedError # pragma: no cover
def export_item(self, item: ImageItem):
raise NotImplementedError # pragma: no cover
def post_export(self):
raise NotImplementedError # pragma: no cover
def export_from(self, items: Iterator[ImageItem]):
self.pre_export()
names = get_task_names()
if names:
desc = f'{self.__class__.__name__} - {".".join(names)}'
else:
desc = f'{self.__class__.__name__}'
for item in tqdm(items, desc=desc):
self.export_item(item)
self.post_export()
def reset(self):
raise NotImplementedError # pragma: no cover
class LocalDirectoryExporter(BaseExporter):
def __init__(self, output_dir, clear: bool = False):
self.output_dir = output_dir
self.clear = clear
def pre_export(self):
if self.clear and os.path.exists(self.output_dir):
remove(self.output_dir)
os.makedirs(self.output_dir, exist_ok=True)
def export_item(self, item: ImageItem):
raise NotImplementedError # pragma: no cover
def post_export(self):
pass
def reset(self):
raise NotImplementedError # pragma: no cover
class SaveExporter(LocalDirectoryExporter):
def __init__(self, output_dir, clear: bool = False, no_meta: bool = False,
skip_when_image_exist: bool = False):
LocalDirectoryExporter.__init__(self, output_dir, clear)
self.no_meta = no_meta
self.untitles = 0
self.skip_when_image_exist = skip_when_image_exist
def export_item(self, item: ImageItem):
if 'filename' in item.meta:
filename = item.meta['filename']
else:
self.untitles += 1
filename = f'untited_{self.untitles}.png'
full_filename = os.path.join(self.output_dir, filename)
full_directory = os.path.dirname(full_filename)
if full_directory:
os.makedirs(full_directory, exist_ok=True)
item.save(full_filename, no_meta=self.no_meta, skip_when_image_exist=self.skip_when_image_exist)
def reset(self):
self.untitles = 0
|