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