From 95dcd640546588928655c9b6f8bf6ffb2c1b80d5 Mon Sep 17 00:00:00 2001 From: narugo1992 Date: Mon, 2 Sep 2024 17:08:34 +0800 Subject: [PATCH] dev(narugo): add YandeWebpDataPool --- cheesechaser/datapool/__init__.py | 8 ++-- cheesechaser/datapool/yande.py | 64 ++++++++++++++++++++++---- docs/source/api_doc/datapool/yande.rst | 8 ++++ 3 files changed, 69 insertions(+), 11 deletions(-) diff --git a/cheesechaser/datapool/__init__.py b/cheesechaser/datapool/__init__.py index 7a3f5bdc8..9a17c490f 100644 --- a/cheesechaser/datapool/__init__.py +++ b/cheesechaser/datapool/__init__.py @@ -1,8 +1,10 @@ from .anime_pictures import AnimePicturesDataPool from .bangumibase import BangumiBaseDataPool -from .base import DataLocation, DataPool, HfBasedDataPool, IncrementIDDataPool, InvalidResourceDataError, FileUnrecognizableError, ResourceNotFoundError +from .base import DataLocation, DataPool, HfBasedDataPool, IncrementIDDataPool, InvalidResourceDataError, \ + FileUnrecognizableError, ResourceNotFoundError from .civitai import CivitaiDataPool -from .danbooru import DanbooruDataPool, DanbooruStableDataPool, DanbooruNewestDataPool, DanbooruWebpDataPool, DanbooruNewestWebpDataPool +from .danbooru import DanbooruDataPool, DanbooruStableDataPool, DanbooruNewestDataPool, DanbooruWebpDataPool, \ + DanbooruNewestWebpDataPool from .fancaps import FancapsDataPool from .gelbooru import GelbooruDataPool, GelbooruWebpDataPool from .hentaicosplay import HentaiCosplayDataPool @@ -12,5 +14,5 @@ from .realbooru import RealbooruDataPool from .table import TableBasedHfDataPool, SimpleTableHfDataPool from .threedbooru import ThreedbooruDataPool -from .yande import YandeDataPool +from .yande import YandeDataPool, YandeWebpDataPool from .zerochan import ZerochanWebpDataPool, ZerochanDataPool diff --git a/cheesechaser/datapool/yande.py b/cheesechaser/datapool/yande.py index 99360960c..da5de7030 100644 --- a/cheesechaser/datapool/yande.py +++ b/cheesechaser/datapool/yande.py @@ -1,12 +1,17 @@ """ -This module provides a data pool implementation for Yande image data. +This module provides data pool implementations for Yande image data. It extends the IncrementIDDataPool class to specifically handle Yande image data, -utilizing a predefined Hugging Face repository for data storage and indexing. +utilizing predefined Hugging Face repositories for data storage and indexing. + +The module includes two main classes: +1. YandeDataPool: For managing original Yande image data. +2. YandeWebpDataPool: For managing WebP-formatted Yande image data. .. note:: - The dataset `deepghs/yande_full `_ - is gated, you have to get the access of it before using this module. + The datasets `deepghs/yande_full `_ and + `deepghs/yande-webp-4Mpixel `_ + are gated, you have to get access to them before using this module. """ from typing import Optional @@ -35,10 +40,6 @@ class YandeDataPool(IncrementIDDataPool): Usage: >>> yande_pool = YandeDataPool() >>> yande_pool_with_token = YandeDataPool(hf_token='your_token_here') - - .. note:: - This class assumes that both data and index information are stored - in the same repository (_YANDE_REPO). """ def __init__(self, revision: str = 'main', hf_token: Optional[str] = None): @@ -62,3 +63,50 @@ def __init__(self, revision: str = 'main', hf_token: Optional[str] = None): idx_revision=revision, hf_token=hf_token, ) + + +_YANDE_WEBP_REPO = 'deepghs/yande-webp-4Mpixel' + + +class YandeWebpDataPool(IncrementIDDataPool): + """ + A data pool class for managing WebP-formatted Yande image data. + + This class extends IncrementIDDataPool to provide a specialized implementation + for handling WebP-formatted Yande image data. It uses a predefined Hugging Face + repository for both data storage and indexing. + + :param revision: The revision of the data to use, defaults to 'main'. + :type revision: str + :param hf_token: Optional Hugging Face authentication token. + :type hf_token: Optional[str] + + :ivar data_repo_id: The Hugging Face repository ID for data storage. + :ivar idx_repo_id: The Hugging Face repository ID for indexing. + + Usage: + >>> yande_webp_pool = YandeWebpDataPool() + >>> yande_webp_pool_with_token = YandeWebpDataPool(hf_token='your_token_here') + """ + + def __init__(self, revision: str = 'main', hf_token: Optional[str] = None): + """ + Initialize the YandeWebpDataPool. + + :param revision: The revision of the data to use, defaults to 'main'. + :type revision: str + :param hf_token: Optional Hugging Face authentication token for accessing private repositories. + :type hf_token: Optional[str] + + This constructor initializes the YandeWebpDataPool by calling the parent class constructor + with specific parameters for the WebP-formatted Yande data repository. It sets up both the data + and index repositories to use the same Hugging Face repository and revision. + """ + IncrementIDDataPool.__init__( + self, + data_repo_id=_YANDE_WEBP_REPO, + data_revision=revision, + idx_repo_id=_YANDE_WEBP_REPO, + idx_revision=revision, + hf_token=hf_token, + ) diff --git a/docs/source/api_doc/datapool/yande.rst b/docs/source/api_doc/datapool/yande.rst index 7b6aa2f0d..57500f0e9 100644 --- a/docs/source/api_doc/datapool/yande.rst +++ b/docs/source/api_doc/datapool/yande.rst @@ -14,3 +14,11 @@ YandeDataPool +YandeWebpDataPool +----------------------------------------------------- + +.. autoclass:: YandeWebpDataPool + :members: __doc__,__init__,__module__ + + +