dev(narugo): add anime pictures webp source

deepghs · Sep 9, 2024 · 03d097f · 03d097f
1 parent 2b5b29c
commit 03d097f
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 3 deletions.
diff --git a/cheesechaser/datapool/__init__.py b/cheesechaser/datapool/__init__.py
@@ -1,4 +1,4 @@
-from .anime_pictures import AnimePicturesDataPool
+from .anime_pictures import AnimePicturesDataPool, AnimePicturesWebpDataPool
 from .bangumibase import BangumiBaseDataPool
 from .base import DataLocation, DataPool, HfBasedDataPool, IncrementIDDataPool, InvalidResourceDataError, \
     FileUnrecognizableError, ResourceNotFoundError

diff --git a/cheesechaser/datapool/anime_pictures.py b/cheesechaser/datapool/anime_pictures.py
@@ -32,7 +32,7 @@ class AnimePicturesDataPool(IncrementIDDataPool):
 
     Usage:
         >>> pool = AnimePicturesDataPool()
-        >>> pool = AnimePicturesDataPool(revision='v1.0')
+        >>> pool = AnimePicturesDataPool(revision='main')
         >>> pool = AnimePicturesDataPool(revision='main', hf_token='your_hf_token')
 
     .. note::
@@ -59,3 +59,18 @@ def __init__(self, revision: str = 'main', hf_token: Optional[str] = None):
             idx_revision=revision,
             hf_token=hf_token,
         )
+
+
+_ANIME_PICTURES_WEBP_REPO = 'deepghs/anime_pictures-webp-4Mpixel'
+
+
+class AnimePicturesWebpDataPool(IncrementIDDataPool):
+    def __init__(self, revision: str = 'main', hf_token: Optional[str] = None):
+        IncrementIDDataPool.__init__(
+            self,
+            data_repo_id=_ANIME_PICTURES_WEBP_REPO,
+            data_revision=revision,
+            idx_repo_id=_ANIME_PICTURES_WEBP_REPO,
+            idx_revision=revision,
+            hf_token=hf_token,
+        )
diff --git a/docs/source/api_doc/datapool/anime_pictures.rst b/docs/source/api_doc/datapool/anime_pictures.rst
@@ -14,3 +14,11 @@ AnimePicturesDataPool
 
 
 
+AnimePicturesWebpDataPool
+-----------------------------------------------------
+
+.. autoclass:: AnimePicturesWebpDataPool
+    :members: __doc__,__init__,__module__
+
+
+
diff --git a/test/datapool/test_anime_pictures.py b/test/datapool/test_anime_pictures.py
@@ -1,7 +1,7 @@
 import pytest
 from hbutils.testing import isolated_directory
 
-from cheesechaser.datapool import AnimePicturesDataPool
+from cheesechaser.datapool import AnimePicturesDataPool, AnimePicturesWebpDataPool
 from ..testings import get_testfile, dir_compare
 
 
@@ -16,3 +16,13 @@ def test_anime_pictures_origin(self):
             )
 
             dir_compare('.', get_testfile('anime_pictures_2'))
+
+    def test_anime_pictures_webp(self):
+        with isolated_directory():
+            pool = AnimePicturesWebpDataPool()
+            pool.batch_download_to_directory(
+                resource_ids=range(200000, 200005),
+                dst_dir='.',
+            )
+
+            dir_compare('.', get_testfile('anime_pictures_webp_2'))
diff --git a/test/testfile/anime_pictures_webp_2/200002.webp b/test/testfile/anime_pictures_webp_2/200002.webp
diff --git a/test/testfile/anime_pictures_webp_2/200004.webp b/test/testfile/anime_pictures_webp_2/200004.webp