From 7fc91385078ce9af3503c92f519f11c5b3b8a66f Mon Sep 17 00:00:00 2001
From: skier233 <39396856+skier233@users.noreply.github.com>
Date: Wed, 5 Jun 2024 05:26:00 -0400
Subject: [PATCH] [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
---
 plugins/AITagger/ai_tagger.py  | 21 ++++++++++++++++++++-
 plugins/AITagger/ai_tagger.yml |  2 +-
 plugins/AITagger/config.py     |  1 +
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/plugins/AITagger/ai_tagger.py b/plugins/AITagger/ai_tagger.py
index 2cfba40a..48c619f7 100644
--- a/plugins/AITagger/ai_tagger.py
+++ b/plugins/AITagger/ai_tagger.py
@@ -3,6 +3,8 @@
 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)
         
 
 
diff --git a/plugins/AITagger/ai_tagger.yml b/plugins/AITagger/ai_tagger.yml
index fb8ab162..711bdf72 100644
--- a/plugins/AITagger/ai_tagger.yml
+++ b/plugins/AITagger/ai_tagger.yml
@@ -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
diff --git a/plugins/AITagger/config.py b/plugins/AITagger/config.py
index da7bc5ef..0ee0b253 100644
--- a/plugins/AITagger/config.py
+++ b/plugins/AITagger/config.py
@@ -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"