[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:
skier233
2024-06-05 05:26:00 -04:00
committed by GitHub
parent 8d3631a12c
commit 7fc9138507
3 changed files with 22 additions and 2 deletions

View File

@@ -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)

View 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

View File

@@ -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"