From 2f5caf4ed77661c7ca71e8d6d7a51e5d69e989e7 Mon Sep 17 00:00:00 2001
From: skier233 <39396856+skier233@users.noreply.github.com>
Date: Tue, 30 Jul 2024 00:30:03 -0400
Subject: [PATCH] [AI Tagger] Prepare stash plugin to support new model tags
 (#363)

---
 plugins/AITagger/ai_tagger.py       | 10 +++++-----
 plugins/AITagger/ai_tagger.yml      |  2 +-
 plugins/AITagger/ai_video_result.py |  4 ++--
 plugins/AITagger/media_handler.py   |  9 +++++++++
 plugins/AITagger/tag_mappings.csv   | 19 +++++++++++++++++++
 5 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/plugins/AITagger/ai_tagger.py b/plugins/AITagger/ai_tagger.py
index 918e0ce3..0d2f6268 100644
--- a/plugins/AITagger/ai_tagger.py
+++ b/plugins/AITagger/ai_tagger.py
@@ -154,10 +154,10 @@ async def __tag_images(images):
                         log.error(f"Error processing image: {result['error']}")
                         media_handler.add_error_images([id])
                     else:
-                        actions = result['actions']
-                        action_stashtag_ids = media_handler.get_tag_ids(actions)
-                        action_stashtag_ids.append(media_handler.ai_tagged_tag_id)
-                        media_handler.add_tags_to_image(id, action_stashtag_ids)
+                        tags = media_handler.get_all_tags_from_server_result(result)
+                        stashtag_ids = media_handler.get_tag_ids(tags)
+                        stashtag_ids.append(media_handler.ai_tagged_tag_id)
+                        media_handler.add_tags_to_image(id, stashtag_ids)
 
             log.info(f"Tagged {len(imageIds)} images")
             media_handler.remove_tagme_tags_from_images(imageIds)
@@ -166,7 +166,7 @@ async def __tag_images(images):
         except asyncio.TimeoutError as a:
             log.error(f"Timeout processing images: {a}")
         except Exception as e:
-            log.error(f"Failed to process images: {e}")
+            log.error(f"Failed to process images: {e}Stack trace: {traceback.format_exc()}")
             media_handler.add_error_images(imageIds)
             media_handler.remove_tagme_tags_from_images(imageIds)
         finally:
diff --git a/plugins/AITagger/ai_tagger.yml b/plugins/AITagger/ai_tagger.yml
index 34166da5..9ae1836a 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.6
+version: 1.7
 url: https://github.com/stashapp/CommunityScripts/tree/main/plugins/AITagger
 exec:
   - python
diff --git a/plugins/AITagger/ai_video_result.py b/plugins/AITagger/ai_video_result.py
index 26e0d33f..355c8bd1 100644
--- a/plugins/AITagger/ai_video_result.py
+++ b/plugins/AITagger/ai_video_result.py
@@ -145,8 +145,8 @@ def __mutate_server_result_tags(cls, server_result, model_name, frame_interval):
         tags = {}
         for result in server_result:
             frame_index = result["frame_index"]
-            actions = result["actions"]
-            for action in actions:
+            alltags = media_handler.get_all_tags_from_server_result(result)
+            for action in alltags:
                 tag_name, confidence = action
                 if tag_name not in tags:
                     tags[tag_name] = TagData(ai_model_name=model_name, time_frames=[TagTimeFrame(start=frame_index, end=None, confidence=confidence)])
diff --git a/plugins/AITagger/media_handler.py b/plugins/AITagger/media_handler.py
index 7e538e15..ada83e60 100644
--- a/plugins/AITagger/media_handler.py
+++ b/plugins/AITagger/media_handler.py
@@ -42,6 +42,15 @@ def initialize(connection):
 
 # ----------------- Tag Methods -----------------
 
+
+tag_categories = ["actions", "bodyparts"]
+
+def get_all_tags_from_server_result(result):
+    alltags = []
+    for category in tag_categories:
+        alltags.extend(result.get(category, []))
+    return alltags
+
 def get_tag_ids(tag_names):
     return [get_tag_id(tag_name) for tag_name in tag_names]
 
diff --git a/plugins/AITagger/tag_mappings.csv b/plugins/AITagger/tag_mappings.csv
index a984ba1d..f65efdcf 100644
--- a/plugins/AITagger/tag_mappings.csv
+++ b/plugins/AITagger/tag_mappings.csv
@@ -35,3 +35,22 @@ Undressing,Undressing_AI,15,6,20s,0.5
 Vaginal Penetration,Vaginal Penetration_AI,15,6,20s,0.5
 Vaginal Fucking,Vaginal Fucking_AI,15,6,20s,0.5
 Vibrating,Vibrating_AI,10,6,20s,0.5
+Ass,Ass_AI,-1,6,20s,0.5
+Asshole,Asshole_AI,-1,6,20s,0.5
+Anal Gape,Anal Gape_AI,10,6,20s,0.5
+Boobs,Boobs_AI,-1,6,20s,0.5
+Cum,Cum_AI,10,6,20s,0.5
+Dick,Dick_AI,-1,6,20s,0.5
+Face,Face_AI,-1,6,20s,0.5
+Feet,Feet_AI,-1,6,20s,0.5
+Fingers,Fingers_AI,-1,6,20s,0.5
+Belly Button,Belly Button_AI,-1,6,20s,0.5
+Nipples,Nipples_AI,-1,6,20s,0.5
+Thighs,Thighs_AI,-1,6,20s,0.5
+Lower Legs,Lower Legs_AI,-1,6,20s,0.5
+Tongue,Tongue_AI,10,6,20s,0.5
+Pussy,Pussy_AI,-1,6,20s,0.5
+Pussy Gape,Pussy Gape_AI,10,6,20s,0.5
+Spit,Spit_AI,10,6,20s,0.5
+Oiled,Oiled_AI,10,6,20s,0.5
+Wet (Water),Wet (Water)_AI,10,6,20s,0.5