waifuc/export/base.py

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