mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 15:41:00 +02:00
Vladimir Mandic
93 lines
2.6 KiB
Python
Executable File
93 lines
2.6 KiB
Python
Executable File
#!/bin/env python
|
|
"""
|
|
use clip to interrogate image(s)
|
|
"""
|
|
|
|
import io
|
|
import base64
|
|
import sys
|
|
import os
|
|
import asyncio
|
|
|
|
import filetype
|
|
from PIL import Image
|
|
|
|
from util import log, Map
|
|
import sdapi as sdapi
|
|
|
|
|
|
def decode(encoding):
|
|
if encoding.startswith("data:image/"):
|
|
encoding = encoding.split(";")[1].split(",")[1]
|
|
return Image.open(io.BytesIO(base64.b64decode(encoding)))
|
|
|
|
|
|
def encode(f):
|
|
image = Image.open(f)
|
|
exif = image.getexif()
|
|
with io.BytesIO() as stream:
|
|
image.save(stream, 'JPEG', exif = exif)
|
|
values = stream.getvalue()
|
|
encoded = base64.b64encode(values).decode()
|
|
return encoded
|
|
|
|
|
|
async def interrogate(f):
|
|
if not filetype.is_image(f):
|
|
log.info({ 'interrogate skip': f })
|
|
return
|
|
json = Map({ 'image': encode(f) })
|
|
log.info({ 'interrogate': f })
|
|
# run clip
|
|
json.model = 'clip'
|
|
res = await sdapi.post('/sdapi/v1/interrogate', json)
|
|
# res = sdapi.postsync('/sdapi/v1/interrogate', json)
|
|
caption = ""
|
|
style = ""
|
|
if 'caption' in res:
|
|
caption = res.caption
|
|
log.info({ 'interrogate caption': caption })
|
|
if ', by' in caption:
|
|
style = caption.split(', by')[1].strip()
|
|
log.info({ 'interrogate style': style })
|
|
else:
|
|
log.error({ 'interrogate clip error': res })
|
|
# run booru
|
|
json.model = 'deepdanbooru'
|
|
res = await sdapi.post('/sdapi/v1/interrogate', json)
|
|
# res = sdapi.postsync('/sdapi/v1/interrogate', json)
|
|
keywords = {}
|
|
if 'caption' in res:
|
|
for term in res.caption.split(', '):
|
|
term = term.replace('(', '').replace(')', '').split(':')
|
|
keywords[term[0]] = term[1]
|
|
keywords = dict(sorted(keywords.items(), key=lambda x:x[1], reverse=True))
|
|
log.info({ 'interrogate keywords': keywords })
|
|
else:
|
|
log.error({ 'interrogate booru error': res })
|
|
return caption, keywords, style
|
|
|
|
|
|
async def main():
|
|
sys.argv.pop(0)
|
|
await sdapi.session()
|
|
if len(sys.argv) == 0:
|
|
log.error({ 'interrogate': 'no files specified' })
|
|
for arg in sys.argv:
|
|
if os.path.exists(arg):
|
|
if os.path.isfile(arg):
|
|
await interrogate(arg)
|
|
elif os.path.isdir(arg):
|
|
for root, _dirs, files in os.walk(arg):
|
|
for f in files:
|
|
await interrogate(os.path.join(root, f))
|
|
else:
|
|
log.error({ 'interrogate unknown file type': arg })
|
|
else:
|
|
log.error({ 'interrogate file missing': arg })
|
|
await sdapi.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|