391 lines
17 KiB
Python
391 lines
17 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import io
|
|
import base64
|
|
import requests
|
|
import json
|
|
import folder_paths
|
|
import comfy.model_management
|
|
import re
|
|
|
|
class GeminiPoseAnalyzer:
|
|
"""
|
|
Detailed Pose Description Generator for OpenPose Enhancement
|
|
Generates comprehensive pose descriptions for multiple people
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"api_key": ("STRING", {
|
|
"multiline": False,
|
|
"default": "YOUR_GEMINI_API_KEY_HERE"
|
|
}),
|
|
"analysis_depth": (["ultra_detailed", "detailed", "standard"], {
|
|
"default": "detailed"
|
|
}),
|
|
"focus_areas": (["full_body", "upper_body", "hands_focus", "face_focus"], {
|
|
"default": "full_body"
|
|
}),
|
|
"pose_style": (["natural", "dynamic", "static", "artistic"], {
|
|
"default": "natural"
|
|
}),
|
|
"language": ([
|
|
"japanese", "english", "chinese", "korean", "spanish", "french",
|
|
"german", "russian", "italian", "portuguese", "dutch", "swedish",
|
|
"norwegian", "danish", "finnish", "polish", "czech", "hungarian",
|
|
"turkish", "arabic", "hindi", "thai"
|
|
], {
|
|
"default": "japanese"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"existing_prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"forceInput": True
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("english_prompt", "trans_text")
|
|
FUNCTION = "analyze_pose"
|
|
CATEGORY = "image/analysis"
|
|
|
|
def image_to_base64(self, image_tensor):
|
|
"""Convert ComfyUI image tensor to base64 string"""
|
|
if len(image_tensor.shape) == 4:
|
|
image_tensor = image_tensor.squeeze(0)
|
|
|
|
image_np = (image_tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(image_np)
|
|
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format='JPEG', quality=98)
|
|
image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
return image_base64
|
|
|
|
def get_translation_keyword(self, language):
|
|
"""Get translation keyword for each language"""
|
|
keywords = {
|
|
"japanese": "[日本語翻訳]",
|
|
"chinese": "[中文翻译]",
|
|
"korean": "[한국어번역]",
|
|
"spanish": "[Traducción española]",
|
|
"french": "[Traduction française]",
|
|
"german": "[Deutsche Übersetzung]",
|
|
"russian": "[Русский перевод]",
|
|
"italian": "[Traduzione italiana]",
|
|
"portuguese": "[Tradução portuguesa]",
|
|
"dutch": "[Nederlandse vertaling]",
|
|
"swedish": "[Svensk översättning]",
|
|
"norwegian": "[Norsk oversettelse]",
|
|
"danish": "[Dansk oversættelse]",
|
|
"finnish": "[Suomenkielinen käännös]",
|
|
"polish": "[Polskie tłumaczenie]",
|
|
"czech": "[Český překlad]",
|
|
"hungarian": "[Magyar fordítás]",
|
|
"turkish": "[Türkçe çeviri]",
|
|
"arabic": "[الترجمة العربية]",
|
|
"hindi": "[हिंदी अनुवाद]",
|
|
"thai": "[การแปลภาษาไทย]"
|
|
}
|
|
return keywords.get(language, "[Translation]")
|
|
|
|
def get_language_name(self, language):
|
|
"""Get native language name"""
|
|
names = {
|
|
"japanese": "日本語",
|
|
"english": "English",
|
|
"chinese": "中文",
|
|
"korean": "한국어",
|
|
"spanish": "Español",
|
|
"french": "Français",
|
|
"german": "Deutsch",
|
|
"russian": "Русский",
|
|
"italian": "Italiano",
|
|
"portuguese": "Português",
|
|
"dutch": "Nederlands",
|
|
"swedish": "Svenska",
|
|
"norwegian": "Norsk",
|
|
"danish": "Dansk",
|
|
"finnish": "Suomi",
|
|
"polish": "Polski",
|
|
"czech": "Čeština",
|
|
"hungarian": "Magyar",
|
|
"turkish": "Türkçe",
|
|
"arabic": "العربية",
|
|
"hindi": "हिंदी",
|
|
"thai": "ไทย"
|
|
}
|
|
return names.get(language, language.title())
|
|
|
|
def create_detailed_pose_prompt(self, analysis_depth, focus_areas, pose_style, language):
|
|
"""Create detailed pose description prompt for multiple languages"""
|
|
|
|
translation_keyword = self.get_translation_keyword(language)
|
|
lang_name = self.get_language_name(language)
|
|
|
|
# Base universal instruction in English (Gemini understands English best)
|
|
base_instruction = f"""Create detailed pose descriptions for ALL people in this image. This is for OpenPose enhancement, describing each person's body orientation, weight distribution, limb placement, and gaze direction accurately.
|
|
|
|
🔥 **ABSOLUTE RULES** 🔥
|
|
1. Count ALL people accurately (single or multiple)
|
|
2. Clearly describe positional relationships between people
|
|
3. Detail each body part (head, torso, arms, hands, legs, feet) precisely
|
|
4. No introductory text - only pure pose descriptions
|
|
5. Use {translation_keyword} as keyword to separate English and {lang_name} versions
|
|
|
|
📋 **Required Description Elements**
|
|
|
|
【Overall Composition】
|
|
- Total number of people and their relative positions
|
|
- Basic orientation and posture of each person
|
|
|
|
【Individual Details (per person)】
|
|
- Body direction: front facing, three-quarter view, side profile, back view
|
|
- Weight distribution: weight on left/right leg, balanced stance
|
|
- Head: head position, gaze direction, head tilt
|
|
- Torso: torso angle, chest direction, spine curve
|
|
- Right arm: shoulder position, elbow bend, hand placement
|
|
- Left arm: shoulder position, elbow bend, hand placement
|
|
- Right leg: leg position, knee bend, foot placement
|
|
- Left leg: leg position, knee bend, foot placement
|
|
- Hand states: finger position, hand gestures, touching points
|
|
|
|
===OUTPUT FORMAT===
|
|
|
|
Single person:
|
|
[Detailed English pose description sentence]
|
|
{translation_keyword}[Detailed {lang_name} pose description sentence]
|
|
|
|
Multiple people:
|
|
[Number] people positioned [positional relationship],
|
|
left position: [detailed English pose description],
|
|
center position: [detailed English pose description],
|
|
right position: [detailed English pose description]
|
|
{translation_keyword}[Number] people positioned [positional relationship in {lang_name}], left position: [detailed {lang_name} description], center position: [detailed {lang_name} description], right position: [detailed {lang_name} description]
|
|
|
|
**Important**:
|
|
- Use "left position", "center position", "right position" for positioning
|
|
- Distinguish left/right clearly: "right arm", "left arm", "right leg", "left leg"
|
|
- Include specific angles and distances (slightly, moderately, fully, etc.)
|
|
- Do NOT use "===" markers
|
|
- No explanatory text - only pure descriptive sentences
|
|
- NEVER use words like figure, mannequin, doll that suggest artificial forms
|
|
- Translate naturally and accurately to {lang_name}"""
|
|
|
|
return base_instruction
|
|
|
|
def call_gemini_api(self, api_key, image_base64, prompt):
|
|
"""Call Gemini API for detailed pose analysis"""
|
|
url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={api_key}"
|
|
|
|
payload = {
|
|
"contents": [
|
|
{
|
|
"parts": [
|
|
{
|
|
"text": prompt
|
|
},
|
|
{
|
|
"inline_data": {
|
|
"mime_type": "image/jpeg",
|
|
"data": image_base64
|
|
}
|
|
}
|
|
]
|
|
}
|
|
],
|
|
"generationConfig": {
|
|
"temperature": 0.1, # Low for consistent descriptions
|
|
"topK": 40, # Allow more vocabulary for descriptions
|
|
"topP": 0.8, # Good balance for detailed text
|
|
"maxOutputTokens": 2000, # More tokens for detailed descriptions
|
|
},
|
|
"safetySettings": [
|
|
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_ONLY_HIGH"},
|
|
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
|
|
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
|
|
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_ONLY_HIGH"}
|
|
]
|
|
}
|
|
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
try:
|
|
response = requests.post(url, json=payload, headers=headers, timeout=60)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
|
|
if 'candidates' in result and len(result['candidates']) > 0:
|
|
content = result['candidates'][0]['content']['parts'][0]['text']
|
|
return content.strip()
|
|
else:
|
|
return "pose analysis failed - no candidates"
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"🚨 Gemini API Error: {str(e)}")
|
|
return "API connection error"
|
|
except Exception as e:
|
|
print(f"🚨 Error processing Gemini response: {str(e)}")
|
|
return "processing error"
|
|
|
|
def process_pose_description(self, raw_output):
|
|
"""Process the detailed pose description"""
|
|
print(f"🔍 Raw API output: {raw_output}")
|
|
|
|
# Remove any introductory text
|
|
cleaned = self.remove_introductions(raw_output)
|
|
print(f"🧹 After intro removal: {cleaned}")
|
|
|
|
return cleaned.strip()
|
|
|
|
def remove_introductions(self, text):
|
|
"""Remove introduction patterns from text"""
|
|
intro_patterns = [
|
|
r'^.*?(?:この画像|This image|Looking at|Based on|Analyzing|I can see|The image shows).*?\n',
|
|
r'^.*?(?:分析|analysis|shows|contains|depicts|features).*?\n',
|
|
r'^\*\*.*?\*\*\s*\n',
|
|
r'^#{1,6}\s.*?\n',
|
|
r'^.*?(?:人物|people|person|figures?|mannequins?).*?(?:います|present|visible|shown).*?\n',
|
|
r'^.*?(?:Here|これは|以下).*?\n',
|
|
]
|
|
|
|
for pattern in intro_patterns:
|
|
text = re.sub(pattern, '', text, flags=re.IGNORECASE | re.MULTILINE)
|
|
|
|
return text.strip()
|
|
|
|
def extract_english_description(self, pose_description):
|
|
"""Extract English description for Stable Diffusion"""
|
|
if '[日本語翻訳]' in pose_description:
|
|
english_part = pose_description.split('[日本語翻訳]')[0].strip()
|
|
elif '[Japanese]' in pose_description:
|
|
english_part = pose_description.split('[Japanese]')[0].strip()
|
|
else:
|
|
# If no Japanese marker, check line by line
|
|
lines = pose_description.split('\n')
|
|
english_lines = []
|
|
for line in lines:
|
|
line = line.strip()
|
|
if line and not self.contains_japanese(line):
|
|
english_lines.append(line)
|
|
english_part = ' '.join(english_lines)
|
|
|
|
return english_part
|
|
|
|
def contains_japanese(self, text):
|
|
"""Check if text contains Japanese characters"""
|
|
return bool(re.search(r'[ひらがなカタカナ漢字]', text))
|
|
|
|
def extract_translated_description(self, pose_description, language):
|
|
"""Extract translated description for viewer display"""
|
|
translation_keyword = self.get_translation_keyword(language)
|
|
|
|
if translation_keyword in pose_description:
|
|
translated_part = pose_description.split(translation_keyword)[1].strip()
|
|
else:
|
|
# Fallback: try to find text in the target language
|
|
lines = pose_description.split('\n')
|
|
translated_lines = []
|
|
for line in lines:
|
|
line = line.strip()
|
|
if line and self.contains_target_language(line, language):
|
|
translated_lines.append(line)
|
|
translated_part = '\n'.join(translated_lines) if translated_lines else f"翻訳テキストが見つかりません / Translation not found"
|
|
|
|
return translated_part
|
|
|
|
def contains_target_language(self, text, language):
|
|
"""Check if text contains target language characters"""
|
|
if language == "japanese":
|
|
return bool(re.search(r'[ひらがなカタカナ漢字]', text))
|
|
elif language == "chinese":
|
|
return bool(re.search(r'[\u4e00-\u9fff]', text))
|
|
elif language == "korean":
|
|
return bool(re.search(r'[\uac00-\ud7af]', text))
|
|
elif language == "arabic":
|
|
return bool(re.search(r'[\u0600-\u06ff]', text))
|
|
elif language == "hindi":
|
|
return bool(re.search(r'[\u0900-\u097f]', text))
|
|
elif language == "thai":
|
|
return bool(re.search(r'[\u0e00-\u0e7f]', text))
|
|
elif language == "russian":
|
|
return bool(re.search(r'[\u0400-\u04ff]', text))
|
|
else:
|
|
# For Latin-based languages, check for non-English words
|
|
return len(text) > 0
|
|
|
|
def analyze_pose(self, image, api_key, analysis_depth, focus_areas, pose_style, language, existing_prompt=""):
|
|
"""Main processing: Detailed pose analysis for OpenPose enhancement"""
|
|
if not api_key or api_key == "YOUR_GEMINI_API_KEY_HERE":
|
|
lang_name = self.get_language_name(language)
|
|
return ("Error: Please set your Gemini API key", f"APIキーを設定してください / Please set API key ({lang_name})")
|
|
|
|
try:
|
|
print(f"🚀 Starting detailed pose analysis...")
|
|
|
|
# Convert image to base64
|
|
image_base64 = self.image_to_base64(image)
|
|
|
|
# Create detailed analysis prompt
|
|
analysis_prompt = self.create_detailed_pose_prompt(analysis_depth, focus_areas, pose_style, language)
|
|
print(f"📝 Detailed prompt created")
|
|
|
|
# Call Gemini API
|
|
raw_output = self.call_gemini_api(api_key, image_base64, analysis_prompt)
|
|
print(f"🤖 API response received")
|
|
|
|
# Process the detailed description
|
|
pose_description = self.process_pose_description(raw_output)
|
|
|
|
# Extract English description for output
|
|
english_description = self.extract_english_description(pose_description)
|
|
|
|
# Extract translated description for viewer
|
|
translated_description = self.extract_translated_description(pose_description, language)
|
|
|
|
# Combine with existing prompt
|
|
if existing_prompt.strip():
|
|
final_prompt = f"{existing_prompt.strip()}, {english_description}"
|
|
else:
|
|
final_prompt = english_description
|
|
|
|
print(f"🎯 English Description: {english_description}")
|
|
print(f"🎯 Translated Description ({language}): {translated_description}")
|
|
print(f"🎯 Final Prompt: {final_prompt}")
|
|
|
|
translation_keyword = self.get_translation_keyword(language)
|
|
print(f"🌍 Language: {language}")
|
|
print(f"🔍 Translation keyword: {translation_keyword}")
|
|
print(f"🔍 Contains translation keyword? {translation_keyword in pose_description}")
|
|
print(f"📝 Raw Pose Description: {pose_description}")
|
|
|
|
# 両方を返す（英語は端子出力、翻訳テキストはビューワ用）
|
|
return (final_prompt, translated_description)
|
|
|
|
except Exception as e:
|
|
error_msg = f"🚨 Error in pose analysis: {str(e)}"
|
|
print(error_msg)
|
|
return (error_msg, f"エラー: {error_msg}")
|
|
|
|
# JavaScript不要 - Pythonのみで実装
|
|
|
|
# Node mappings
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GeminiPoseAnalyzer": GeminiPoseAnalyzer
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GeminiPoseAnalyzer": "🌍 Gemini Pose Analyzer [22言語対応]"
|
|
} |