[civitai] implement extracting 'tags' metadata (#8626)
This commit is contained in:
@@ -2090,12 +2090,18 @@ Type
|
||||
Default
|
||||
``false``
|
||||
Example
|
||||
* ``"generation,post,version"``
|
||||
* ``"generation,tags,post,version"``
|
||||
* ``["version", "generation"]``
|
||||
Description
|
||||
Extract additional ``generation``, ``version``, and ``post`` metadata.
|
||||
Extract additional metadata.
|
||||
Supported Values
|
||||
* ``generation``
|
||||
* ``post``
|
||||
* ``tags``
|
||||
* ``version``
|
||||
Note
|
||||
This requires 1 or more additional API requests per image or video.
|
||||
This requires 1 additional API request
|
||||
for each selected value per image or video.
|
||||
|
||||
|
||||
extractor.civitai.nsfw
|
||||
|
||||
@@ -61,13 +61,14 @@ class CivitaiExtractor(Extractor):
|
||||
if isinstance(metadata, str):
|
||||
metadata = metadata.split(",")
|
||||
elif not isinstance(metadata, (list, tuple)):
|
||||
metadata = ("generation", "version", "post")
|
||||
metadata = {"generation", "version", "post", "tags"}
|
||||
self._meta_generation = ("generation" in metadata)
|
||||
self._meta_version = ("version" in metadata)
|
||||
self._meta_post = ("post" in metadata)
|
||||
self._meta_tags = ("tags" in metadata)
|
||||
else:
|
||||
self._meta_generation = self._meta_version = self._meta_post = \
|
||||
False
|
||||
self._meta_tags = False
|
||||
|
||||
def items(self):
|
||||
if models := self.models():
|
||||
@@ -110,8 +111,9 @@ class CivitaiExtractor(Extractor):
|
||||
}
|
||||
|
||||
if self._meta_generation:
|
||||
data["generation"] = \
|
||||
self._extract_meta_generation(file)
|
||||
data["generation"] = self._extract_meta_generation(file)
|
||||
if self._meta_tags:
|
||||
data["tags"] = self._extract_meta_tags(file)
|
||||
if self._meta_version:
|
||||
data["model"], data["version"] = \
|
||||
self._extract_meta_version(file, False)
|
||||
@@ -181,6 +183,8 @@ class CivitaiExtractor(Extractor):
|
||||
file["date"] = self.parse_datetime_iso(file["createdAt"])
|
||||
if self._meta_generation:
|
||||
file["generation"] = self._extract_meta_generation(file)
|
||||
if self._meta_tags:
|
||||
file["tags"] = self._extract_meta_tags(file)
|
||||
yield data
|
||||
|
||||
def _image_reactions(self):
|
||||
@@ -218,6 +222,12 @@ class CivitaiExtractor(Extractor):
|
||||
except Exception as exc:
|
||||
return self.log.traceback(exc)
|
||||
|
||||
def _extract_meta_tags(self, image):
|
||||
try:
|
||||
return self.api.tag_getvotabletags(image["id"])
|
||||
except Exception as exc:
|
||||
return self.log.traceback(exc)
|
||||
|
||||
def _extract_meta_version(self, item, is_post=True):
|
||||
try:
|
||||
if version_id := self._extract_version_id(item, is_post):
|
||||
@@ -821,6 +831,11 @@ class CivitaiTrpcAPI():
|
||||
params = self._type_params(params)
|
||||
return self._pagination(endpoint, params)
|
||||
|
||||
def tag_getvotabletags(self, image_id):
|
||||
endpoint = "tag.getVotableTags"
|
||||
params = {"id": int(image_id), "type": "image"}
|
||||
return self._call(endpoint, params)
|
||||
|
||||
def user(self, username):
|
||||
endpoint = "user.getCreator"
|
||||
params = {"username": username}
|
||||
|
||||
@@ -149,6 +149,19 @@ __tests__ = (
|
||||
"tags": [],
|
||||
"collectionId": None,
|
||||
},
|
||||
"tags[*]": {
|
||||
"automated" : bool,
|
||||
"concrete" : bool,
|
||||
"downVotes" : int,
|
||||
"id" : int,
|
||||
"lastUpvote" : None,
|
||||
"name" : str,
|
||||
"needsReview": bool,
|
||||
"nsfwLevel" : 1,
|
||||
"score" : int,
|
||||
"type" : {"Label", "UserGenerated"},
|
||||
"upVotes" : int,
|
||||
},
|
||||
"model": {
|
||||
"id": 703211,
|
||||
"name": "メイド クラシック/maid classic",
|
||||
|
||||
Reference in New Issue
Block a user