mirror of
https://github.com/stashapp/CommunityScripts.git
synced 2026-05-01 10:16:15 -05:00
[AI Tagger] Add support for zipped images (#324)
* Fix mintags duration being lower than the frame interval * cleanup debug message * add support for zipped images
This commit is contained in:
@@ -3,6 +3,8 @@ import sys
|
||||
import json
|
||||
import subprocess
|
||||
import csv
|
||||
import zipfile
|
||||
import shutil
|
||||
from typing import Any
|
||||
|
||||
# ----------------- Setup -----------------
|
||||
@@ -156,7 +158,19 @@ async def __tag_images(images):
|
||||
async with semaphore:
|
||||
imagePaths = [image['files'][0]['path'] for image in images]
|
||||
imageIds = [image['id'] for image in images]
|
||||
#TODO
|
||||
|
||||
temp_files = []
|
||||
for i, path in enumerate(imagePaths):
|
||||
if '.zip' in path:
|
||||
zip_index = path.index('.zip') + 4
|
||||
zip_path, img_path = path[:zip_index], path[zip_index+1:].replace('\\', '/')
|
||||
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
|
||||
temp_path = os.path.join(config.temp_image_dir, img_path)
|
||||
os.makedirs(os.path.dirname(temp_path), exist_ok=True)
|
||||
zip_ref.extract(img_path, config.temp_image_dir)
|
||||
imagePaths[i] = os.path.abspath(os.path.normpath(temp_path))
|
||||
temp_files.append(imagePaths[i])
|
||||
|
||||
try:
|
||||
server_results = ImageResult(**await process_images_async(imagePaths))
|
||||
process_server_image_results(server_results, imageIds)
|
||||
@@ -166,6 +180,11 @@ async def __tag_images(images):
|
||||
stash.update_images({"ids": imageIds, "tag_ids": {"ids": [tagme_tag_id], "mode": "REMOVE"}})
|
||||
finally:
|
||||
increment_progress()
|
||||
for temp_file in temp_files:
|
||||
if os.path.isdir(temp_file):
|
||||
shutil.rmtree(temp_file)
|
||||
else:
|
||||
os.remove(temp_file)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
name: AI Tagger
|
||||
description: Tag videos and Images with Locally hosted AI using Skier's Patreon AI models
|
||||
version: 1.1
|
||||
version: 1.2
|
||||
url: https://github.com/stashapp/CommunityScripts/tree/main/plugins/AITagger
|
||||
exec:
|
||||
- python
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
API_BASE_URL = 'http://localhost:8000'
|
||||
IMAGE_REQUEST_BATCH_SIZE = 320
|
||||
CONCURRENT_TASK_LIMIT = 10
|
||||
temp_image_dir = "./temp_images"
|
||||
|
||||
ai_base_tag_name = "AI"
|
||||
tagme_tag_name = "AI_TagMe"
|
||||
|
||||
Reference in New Issue
Block a user