diff --git a/modules/interrogate/deepseek.py b/modules/interrogate/deepseek.py index b2d340248..8611efb82 100644 --- a/modules/interrogate/deepseek.py +++ b/modules/interrogate/deepseek.py @@ -18,32 +18,31 @@ from modules import shared, devices, paths, sd_models # model_path = "deepseek-ai/deepseek-vl2-small" vl_gpt = None vl_chat_processor = None +loaded_repo = None class fake_attrdict(): - class AttrDict(dict): # dot notation access to dictionary attributes + class AttrDict(dict): # dot notation access to dictionary attributes __getattr__ = dict.get __setattr__ = dict.__setitem__ __delattr__ = dict.__delitem__ -# def fake_is_flash_attn_2_available(): -# return False - -def predict(question, image, repo): - global vl_gpt, vl_chat_processor # pylint: disable=global-statement +def load(repo: str): + """Load DeepSeek VL2 model (experimental).""" + global vl_gpt, vl_chat_processor, loaded_repo # pylint: disable=global-statement if not shared.cmd_opts.experimental: shared.log.error(f'Interrogate: type=vlm model="DeepSeek VL2" repo="{repo}" is experimental-only') - return '' + return False folder = os.path.join(paths.script_path, 'repositories', 'deepseek-vl2') if not os.path.exists(folder): shared.log.error(f'Interrogate: type=vlm model="DeepSeek VL2" repo="{repo}" deepseek-vl2 repo not found') - return '' - if vl_gpt is None: + return False + if vl_gpt is None or loaded_repo != repo: sys.modules['attrdict'] = fake_attrdict from transformers.models.llama import modeling_llama modeling_llama.LlamaFlashAttention2 = modeling_llama.LlamaAttention - _deekseek_vl = importlib.import_module('repositories.deepseek-vl2.deepseek_vl2') + importlib.import_module('repositories.deepseek-vl2.deepseek_vl2') deekseek_vl_models = importlib.import_module('repositories.deepseek-vl2.deepseek_vl2.models') vl_chat_processor = deekseek_vl_models.DeepseekVLV2Processor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir) vl_gpt = AutoModelForCausalLM.from_pretrained( @@ -51,7 +50,31 @@ def predict(question, image, repo): trust_remote_code=True, cache_dir=shared.opts.hfcache_dir, ) - vl_gpt = vl_gpt.to(device=devices.device, dtype=devices.dtype).eval() + vl_gpt.to(dtype=devices.dtype) + vl_gpt.eval() + loaded_repo = repo + shared.log.info(f'Interrogate: type=vlm model="DeepSeek VL2" repo="{repo}"') + sd_models.move_model(vl_gpt, devices.device) + return True + + +def unload(): + """Release DeepSeek VL2 model from GPU/memory.""" + global vl_gpt, vl_chat_processor, loaded_repo # pylint: disable=global-statement + if vl_gpt is not None: + shared.log.debug(f'DeepSeek unload: model="{loaded_repo}"') + sd_models.move_model(vl_gpt, devices.cpu, force=True) + vl_gpt = None + vl_chat_processor = None + loaded_repo = None + devices.torch_gc(force=True) + else: + shared.log.debug('DeepSeek unload: no model loaded') + + +def predict(question, image, repo): + if not load(repo): + return '' if len(question) < 2: question = "Describe the image." diff --git a/modules/interrogate/joycaption.py b/modules/interrogate/joycaption.py index 114888f4e..c8d445d9e 100644 --- a/modules/interrogate/joycaption.py +++ b/modules/interrogate/joycaption.py @@ -57,26 +57,43 @@ llava_model: LlavaForConditionalGeneration = None opts = JoyOptions() -@torch.no_grad() -def predict(question: str, image, vqa_model: str = None) -> str: - global llava_model, processor # pylint: disable=global-statement - opts.max_new_tokens = shared.opts.interrogate_vlm_max_length - if vqa_model is not None and opts.repo != vqa_model: - opts.repo = vqa_model +def load(repo: str = None): + """Load JoyCaption model.""" + global llava_model, processor # pylint: disable=global-statement + repo = repo or opts.repo + if llava_model is None or opts.repo != repo: + opts.repo = repo llava_model = None - if llava_model is None: shared.log.info(f'Interrogate: type=vlm model="JoyCaption" {str(opts)}') - - processor = AutoProcessor.from_pretrained(opts.repo, max_pixels=1024*1024) + processor = AutoProcessor.from_pretrained(repo, max_pixels=1024*1024, cache_dir=shared.opts.hfcache_dir) quant_args = model_quant.create_config(module='LLM') llava_model = LlavaForConditionalGeneration.from_pretrained( - opts.repo, + repo, torch_dtype=devices.dtype, - device_map="auto", cache_dir=shared.opts.hfcache_dir, **quant_args, ) llava_model.eval() + sd_models.move_model(llava_model, devices.device) + + +def unload(): + """Release JoyCaption model from GPU/memory.""" + global llava_model, processor # pylint: disable=global-statement + if llava_model is not None: + shared.log.debug(f'JoyCaption unload: model="{opts.repo}"') + sd_models.move_model(llava_model, devices.cpu, force=True) + llava_model = None + processor = None + devices.torch_gc(force=True) + else: + shared.log.debug('JoyCaption unload: no model loaded') + + +@torch.no_grad() +def predict(question: str, image, vqa_model: str = None) -> str: + opts.max_new_tokens = shared.opts.interrogate_vlm_max_length + load(vqa_model) if len(question) < 2: question = "Describe the image." @@ -86,9 +103,8 @@ def predict(question: str, image, vqa_model: str = None) -> str: { "role": "user", "content": question }, ] convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True) - inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(devices.device) # Process the inputs + inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(devices.device) inputs['pixel_values'] = inputs['pixel_values'].to(devices.dtype) - sd_models.move_model(llava_model, devices.device) with devices.inference_context(): generate_ids = llava_model.generate( # Generate the captions **inputs, diff --git a/modules/interrogate/joytag.py b/modules/interrogate/joytag.py index ae48cb9c6..f78e49ae5 100644 --- a/modules/interrogate/joytag.py +++ b/modules/interrogate/joytag.py @@ -16,7 +16,7 @@ import torchvision.transforms.functional as TVF import einops from einops.layers.torch import Rearrange import huggingface_hub -from modules import shared, devices +from modules import shared, devices, sd_models model = None @@ -1034,22 +1034,39 @@ def prepare_image(image: Image.Image, target_size: int) -> torch.Tensor: return image_tensor -def predict(image: Image.Image): - global model, tags # pylint: disable=global-statement +def load(): + """Load JoyTag model.""" + global model, tags # pylint: disable=global-statement if model is None: folder = huggingface_hub.snapshot_download(MODEL_REPO, cache_dir=shared.opts.hfcache_dir) model = VisionModel.load_model(folder) - model = model.to(device=devices.device, dtype=devices.dtype) + model.to(dtype=devices.dtype) model.eval() with open(os.path.join(folder, 'top_tags.txt'), 'r', encoding='utf8') as f: tags = [line.strip() for line in f.readlines() if line.strip()] - shared.log.info(f'Interrogate: type=vlm model="JoyCaption" repo="{MODEL_REPO}" tags={len(tags)}') + shared.log.info(f'Interrogate: type=vlm model="JoyTag" repo="{MODEL_REPO}" tags={len(tags)}') + sd_models.move_model(model, devices.device) + + +def unload(): + """Release JoyTag model from GPU/memory.""" + global model, tags # pylint: disable=global-statement + if model is not None: + shared.log.debug('JoyTag unload') + sd_models.move_model(model, devices.cpu, force=True) + model = None + tags = None + devices.torch_gc(force=True) + else: + shared.log.debug('JoyTag unload: no model loaded') + + +def predict(image: Image.Image): + load() image_tensor = prepare_image(image, model.image_size).unsqueeze(0).to(device=devices.device, dtype=devices.dtype) - model = model.to(devices.device) with devices.inference_context(): - preds = model({ 'image': image_tensor }) + preds = model({'image': image_tensor}) tag_preds = preds['tags'].sigmoid().cpu() - model = model.to(devices.cpu) scores = {tags[i]: tag_preds[0][i] for i in range(len(tags))} if shared.opts.interrogate_score: predicted_tags = [f'{tag}:{score:.2f}' for tag, score in scores.items() if score > THRESHOLD] diff --git a/modules/interrogate/moondream3.py b/modules/interrogate/moondream3.py index 739e26f3b..4983cfcd3 100644 --- a/modules/interrogate/moondream3.py +++ b/modules/interrogate/moondream3.py @@ -8,6 +8,7 @@ import torch import transformers from PIL import Image from modules import shared, devices, sd_models +from modules.interrogate import vqa_detection # Debug logging - function-based to avoid circular import @@ -219,32 +220,14 @@ def point(image: Image.Image, object_name: str, repo: str): with devices.inference_context(): result = model.point(image, object_name) - # Debug: Log the actual result to understand the format debug(f'VQA interrogate: handler=moondream3 point_raw_result="{result}" type={type(result)}') if isinstance(result, dict): debug(f'VQA interrogate: handler=moondream3 point_raw_result_keys={list(result.keys())}') - # Parse and validate coordinates - # Handle dict format: {'points': [{'x': 0.733, 'y': 0.442}, {'x': 0.5, 'y': 0.6}, ...]} - if isinstance(result, dict) and 'points' in result: - points_list = result['points'] - if points_list and len(points_list) > 0: - coordinates = [] - for point_data in points_list: # Iterate ALL points - if 'x' in point_data and 'y' in point_data: - x = max(0.0, min(1.0, float(point_data['x']))) - y = max(0.0, min(1.0, float(point_data['y']))) - coordinates.append((x, y)) - if coordinates: - debug(f'VQA interrogate: handler=moondream3 point_result={len(coordinates)} points found') - return coordinates - # Fallback: try simple list/tuple format [x, y] (for compatibility) - elif isinstance(result, (list, tuple)) and len(result) == 2: - x, y = result - x = max(0.0, min(1.0, float(x))) - y = max(0.0, min(1.0, float(y))) - debug('VQA interrogate: handler=moondream3 point_result=1 point found') - return [(x, y)] # Return as list for consistency + points = vqa_detection.parse_points(result) + if points: + debug(f'VQA interrogate: handler=moondream3 point_result={len(points)} points found') + return points debug('VQA interrogate: handler=moondream3 point_result=not found') return None @@ -274,31 +257,11 @@ def detect(image: Image.Image, object_name: str, repo: str, max_objects: int = 1 with devices.inference_context(): result = model.detect(image, object_name) - # Debug: Log the actual result to understand the format debug(f'VQA interrogate: handler=moondream3 detect_raw_result="{result}" type={type(result)}') if isinstance(result, dict): debug(f'VQA interrogate: handler=moondream3 detect_raw_result_keys={list(result.keys())}') - # Parse detections - # Expected format: {'objects': [{'x_min': 0.1, 'y_min': 0.2, 'x_max': 0.5, 'y_max': 0.8}, ...]} - detections = [] - - if isinstance(result, dict) and 'objects' in result: - objects = result['objects'][:max_objects] # Limit to max_objects - for i, obj in enumerate(objects): - if all(k in obj for k in ['x_min', 'y_min', 'x_max', 'y_max']): - bbox = [ - max(0.0, min(1.0, float(obj['x_min']))), - max(0.0, min(1.0, float(obj['y_min']))), - max(0.0, min(1.0, float(obj['x_max']))), - max(0.0, min(1.0, float(obj['y_max']))) - ] - detections.append({ - 'bbox': bbox, - 'label': object_name, - 'confidence': obj.get('confidence', 1.0) # Default confidence if not provided - }) - + detections = vqa_detection.parse_detections(result, object_name, max_objects) debug(f'VQA interrogate: handler=moondream3 detect_result={len(detections)} objects found') return detections @@ -376,55 +339,33 @@ def predict(question: str, image: Image.Image, repo: str, model_name: str = None elif mode == 'point': # Extract object name from question - case insensitive, preserve object names object_name = question - # Remove trigger phrases (case-insensitive) for phrase in ['point at', 'where is', 'locate', 'find']: object_name = re.sub(rf'\b{phrase}\b', '', object_name, flags=re.IGNORECASE) - # Remove punctuation and extra whitespace object_name = re.sub(r'[?.!,]', '', object_name).strip() - # Remove leading "the" only object_name = re.sub(r'^\s*the\s+', '', object_name, flags=re.IGNORECASE) debug(f'VQA interrogate: handler=moondream3 point_extracted_object="{object_name}"') result = point(image, object_name, repo) if result: - # Handle multiple instances - return text and store points for drawing - if len(result) == 1: - text = f"Found at coordinates: ({result[0][0]:.3f}, {result[0][1]:.3f})" - else: - # Multiple instances found - format with count - lines = [f"Found {len(result)} instances:"] - for i, (x, y) in enumerate(result, 1): - lines.append(f" {i}. ({x:.3f}, {y:.3f})") - text = '\n'.join(lines) - # Store detection data on VQA singleton for annotation from modules.interrogate import vqa vqa.get_instance().last_detection_data = {'points': result} - return text + return vqa_detection.format_points_text(result) return "Object not found" elif mode == 'detect': # Extract object name from question - case insensitive object_name = question - # Remove trigger phrases (case-insensitive) for phrase in ['detect', 'find all', 'bounding box', 'bbox', 'find']: object_name = re.sub(rf'\b{phrase}\b', '', object_name, flags=re.IGNORECASE) - # Remove punctuation and extra whitespace object_name = re.sub(r'[?.!,]', '', object_name).strip() - # Remove leading "the" only object_name = re.sub(r'^\s*the\s+', '', object_name, flags=re.IGNORECASE) - # Remove "and" and get first object (model detects one type at a time) if ' and ' in object_name.lower(): object_name = re.split(r'\s+and\s+', object_name, flags=re.IGNORECASE)[0].strip() debug(f'VQA interrogate: handler=moondream3 detect_extracted_object="{object_name}"') results = detect(image, object_name, repo, max_objects=kwargs.get('max_objects', 10)) - # Format as string for display and store detections for drawing if results: - lines = [f"{det['label']}: [{det['bbox'][0]:.3f}, {det['bbox'][1]:.3f}, {det['bbox'][2]:.3f}, {det['bbox'][3]:.3f}] (confidence: {det['confidence']:.2f})" - for det in results] - text = '\n'.join(lines) - # Store detection data on VQA singleton for annotation from modules.interrogate import vqa vqa.get_instance().last_detection_data = {'detections': results} - return text + return vqa_detection.format_detections_text(results) return "No objects detected" else: # mode == 'query' if len(question) < 2: @@ -447,3 +388,17 @@ def clear_cache(): image_cache.clear() debug(f'VQA interrogate: handler=moondream3 cleared image cache cache_size_was={cache_size}') shared.log.debug(f'Moondream3: Cleared image cache ({cache_size} entries)') + + +def unload(): + """Release Moondream 3 model from GPU/memory.""" + global moondream3_model, loaded # pylint: disable=global-statement + if moondream3_model is not None: + shared.log.debug(f'Moondream3 unload: model="{loaded}"') + sd_models.move_model(moondream3_model, devices.cpu, force=True) + moondream3_model = None + loaded = None + clear_cache() + devices.torch_gc(force=True) + else: + shared.log.debug('Moondream3 unload: no model loaded') diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index aa1fad562..02c132a33 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -9,7 +9,7 @@ import transformers import transformers.dynamic_module_utils from PIL import Image from modules import shared, devices, errors, model_quant, sd_models, sd_models_compile, ui_symbols -from modules.interrogate import vqa_draw +from modules.interrogate import vqa_detection # Debug logging - function-based to avoid circular import @@ -417,10 +417,28 @@ class VQA: self._load_sa2(repo) elif 'fastvlm' in repo_lower: self._load_fastvlm(repo) + elif 'moondream3' in repo_lower: + from modules.interrogate import moondream3 + moondream3.load_model(repo) + shared.log.info(f'VQA load: model="{model_name}" loaded (external handler)') + return + elif 'joytag' in repo_lower: + from modules.interrogate import joytag + joytag.load() + shared.log.info(f'VQA load: model="{model_name}" loaded (external handler)') + return + elif 'joycaption' in repo_lower: + from modules.interrogate import joycaption + joycaption.load(repo) + shared.log.info(f'VQA load: model="{model_name}" loaded (external handler)') + return + elif 'deepseek' in repo_lower: + from modules.interrogate import deepseek + deepseek.load(repo) + shared.log.info(f'VQA load: model="{model_name}" loaded (external handler)') + return else: - # Models with external handlers (moondream3, joytag, joycaption, deepseek) - # don't support pre-loading through this method - shared.log.warning(f'VQA load: no pre-loader for model="{model_name}" (external handler)') + shared.log.warning(f'VQA load: no pre-loader for model="{model_name}"') return sd_models.move_model(self.model, devices.device) @@ -1066,19 +1084,10 @@ class VQA: debug(f'VQA interrogate: handler=moondream method=point target="{target}"') result = self.model.point(image, target) debug(f'VQA interrogate: handler=moondream point_raw_result={result}') - # Parse points: {'points': [{'x': 0.5, 'y': 0.5}, ...]} - if isinstance(result, dict) and 'points' in result: - points = [(p['x'], p['y']) for p in result['points'] if 'x' in p and 'y' in p] - if points: - if len(points) == 1: - text = f"Found at: ({points[0][0]:.3f}, {points[0][1]:.3f})" - else: - lines = [f"Found {len(points)} instances:"] - for i, (x, y) in enumerate(points, 1): - lines.append(f" {i}. ({x:.3f}, {y:.3f})") - text = '\n'.join(lines) - self.last_detection_data = {'points': points} - return text + points = vqa_detection.parse_points(result) + if points: + self.last_detection_data = {'points': points} + return vqa_detection.format_points_text(points) return "Object not found" elif question.lower().startswith('detect ') or question == 'DETECT_MODE': target = question[7:].strip() if question.lower().startswith('detect ') else '' @@ -1087,36 +1096,23 @@ class VQA: debug(f'VQA interrogate: handler=moondream method=detect target="{target}"') result = self.model.detect(image, target) debug(f'VQA interrogate: handler=moondream detect_raw_result={result}') - # Parse objects: {'objects': [{'x_min': .1, 'y_min': .2, 'x_max': .5, 'y_max': .8}, ...]} - if isinstance(result, dict) and 'objects' in result: - detections = [] - for obj in result['objects']: - if all(k in obj for k in ['x_min', 'y_min', 'x_max', 'y_max']): - detections.append({ - 'bbox': [obj['x_min'], obj['y_min'], obj['x_max'], obj['y_max']], - 'label': target - }) - if detections: - lines = [f"{d['label']}: [{d['bbox'][0]:.3f}, {d['bbox'][1]:.3f}, {d['bbox'][2]:.3f}, {d['bbox'][3]:.3f}]" for d in detections] - self.last_detection_data = {'detections': detections} - return '\n'.join(lines) + detections = vqa_detection.parse_detections(result, target) + if detections: + self.last_detection_data = {'detections': detections} + return vqa_detection.format_detections_text(detections, include_confidence=False) return "No objects detected" elif question == 'DETECT_GAZE' or question.lower() == 'detect gaze': debug('VQA interrogate: handler=moondream method=detect_gaze') - # First detect faces to get eye regions faces = self.model.detect(image, "face") debug(f'VQA interrogate: handler=moondream detect_gaze faces={faces}') if faces.get('objects'): - face = faces['objects'][0] # Use first face - eye_x = (face['x_min'] + face['x_max']) / 2 - eye_y = face['y_min'] + (face['y_max'] - face['y_min']) * 0.3 # Approximate eye level + eye_x, eye_y = vqa_detection.calculate_eye_position(faces['objects'][0]) result = self.model.detect_gaze(image, eye=(eye_x, eye_y)) debug(f'VQA interrogate: handler=moondream detect_gaze result={result}') if result.get('gaze'): gaze = result['gaze'] - text = f"Gaze direction: ({gaze['x']:.3f}, {gaze['y']:.3f})" self.last_detection_data = {'points': [(gaze['x'], gaze['y'])]} - return text + return f"Gaze direction: ({gaze['x']:.3f}, {gaze['y']:.3f})" return "No face/gaze detected" else: debug(f'VQA interrogate: handler=moondream method=query question="{question}" reasoning={thinking_mode}') @@ -1360,7 +1356,7 @@ class VQA: detections = self.last_detection_data.get('detections', None) points = self.last_detection_data.get('points', None) if detections or points: - self.last_annotated_image = vqa_draw.draw_bounding_boxes(image, detections or [], points) + self.last_annotated_image = vqa_detection.draw_bounding_boxes(image, detections or [], points) debug(f'VQA interrogate: handler={handler} created annotated image detections={len(detections) if detections else 0} points={len(points) if points else 0}') debug(f'VQA interrogate: handler={handler} response_after_clean="{answer}" has_annotation={self.last_annotated_image is not None}') diff --git a/modules/interrogate/vqa_detection.py b/modules/interrogate/vqa_detection.py new file mode 100644 index 000000000..192d669bc --- /dev/null +++ b/modules/interrogate/vqa_detection.py @@ -0,0 +1,207 @@ +# VQA Detection Utilities +# Parsing, formatting, and drawing functions for detection results (points, bboxes, gaze) + +from PIL import Image, ImageDraw, ImageFont +from modules import shared + + +def parse_points(result) -> list: + """Parse and validate point coordinates from model result. + + Args: + result: Model output, typically dict with 'points' key or list of coordinates + + Returns: + List of (x, y) tuples with coordinates clamped to 0-1 range. + """ + points = [] + + # Dict format: {'points': [{'x': 0.5, 'y': 0.5}, ...]} + if isinstance(result, dict) and 'points' in result: + points_list = result['points'] + if points_list and len(points_list) > 0: + for point_data in points_list: + if isinstance(point_data, dict) and 'x' in point_data and 'y' in point_data: + x = max(0.0, min(1.0, float(point_data['x']))) + y = max(0.0, min(1.0, float(point_data['y']))) + points.append((x, y)) + + # Fallback for simple [x, y] format + elif isinstance(result, (list, tuple)) and len(result) == 2: + try: + x = max(0.0, min(1.0, float(result[0]))) + y = max(0.0, min(1.0, float(result[1]))) + points.append((x, y)) + except (ValueError, TypeError): + pass + + return points + + +def parse_detections(result, label: str, max_objects: int = None) -> list: + """Parse and validate detection bboxes from model result. + + Args: + result: Model output, typically dict with 'objects' key + label: Label to assign to detected objects + max_objects: Maximum number of objects to return (None for all) + + Returns: + List of {'bbox': [x1,y1,x2,y2], 'label': str, 'confidence': float} + with coordinates clamped to 0-1 range. + """ + detections = [] + + if isinstance(result, dict) and 'objects' in result: + objects = result['objects'] + if max_objects is not None: + objects = objects[:max_objects] + + for obj in objects: + if all(k in obj for k in ['x_min', 'y_min', 'x_max', 'y_max']): + bbox = [ + max(0.0, min(1.0, float(obj['x_min']))), + max(0.0, min(1.0, float(obj['y_min']))), + max(0.0, min(1.0, float(obj['x_max']))), + max(0.0, min(1.0, float(obj['y_max']))) + ] + detections.append({ + 'bbox': bbox, + 'label': label, + 'confidence': obj.get('confidence', 1.0) + }) + + return detections + + +def format_points_text(points: list) -> str: + """Format point coordinates as human-readable text. + + Args: + points: List of (x, y) tuples with normalized coordinates + + Returns: + Formatted text string describing the points. + """ + if not points: + return "Object not found" + + if len(points) == 1: + return f"Found at: ({points[0][0]:.3f}, {points[0][1]:.3f})" + + lines = [f"Found {len(points)} instances:"] + for i, (x, y) in enumerate(points, 1): + lines.append(f" {i}. ({x:.3f}, {y:.3f})") + return '\n'.join(lines) + + +def format_detections_text(detections: list, include_confidence: bool = True) -> str: + """Format detections with bboxes as human-readable text. + + Args: + detections: List of detection dicts with 'bbox', 'label', 'confidence' + include_confidence: Whether to include confidence scores in output + + Returns: + Formatted text string describing the detections. + """ + if not detections: + return "No objects detected" + + lines = [] + for det in detections: + bbox = det['bbox'] + label = det.get('label', 'object') + confidence = det.get('confidence', 1.0) + + if include_confidence and confidence < 1.0: + lines.append(f"{label}: [{bbox[0]:.3f}, {bbox[1]:.3f}, {bbox[2]:.3f}, {bbox[3]:.3f}] (confidence: {confidence:.2f})") + else: + lines.append(f"{label}: [{bbox[0]:.3f}, {bbox[1]:.3f}, {bbox[2]:.3f}, {bbox[3]:.3f}]") + + return '\n'.join(lines) + + +def calculate_eye_position(face_bbox: dict) -> tuple: + """Calculate approximate eye position from face bounding box. + + Args: + face_bbox: Dict with 'x_min', 'y_min', 'x_max', 'y_max' keys + + Returns: + (eye_x, eye_y) tuple with normalized coordinates. + """ + eye_x = (face_bbox['x_min'] + face_bbox['x_max']) / 2 + eye_y = face_bbox['y_min'] + (face_bbox['y_max'] - face_bbox['y_min']) * 0.3 # Approximate eye level + return (eye_x, eye_y) + + +def draw_bounding_boxes(image: Image.Image, detections: list, points: list = None) -> Image.Image: + """ + Draw bounding boxes and/or points on an image. + + Args: + image: PIL Image to annotate + detections: List of detection dicts with format: + [{'label': str, 'bbox': [x1, y1, x2, y2], 'confidence': float}, ...] + where coordinates are normalized 0-1 + points: Optional list of (x, y) tuples with normalized 0-1 coordinates + + Returns: + Annotated PIL Image with boxes and labels drawn, or None if no annotations + """ + if not detections and not points: + return None + + # Create a copy to avoid modifying original + annotated = image.copy() + draw = ImageDraw.Draw(annotated) + width, height = image.size + + # Try to load a font, fall back to default if unavailable + try: + font_size = max(12, int(min(width, height) * 0.02)) + font_path = shared.opts.font or "javascript/notosans-nerdfont-regular.ttf" + font = ImageFont.truetype(font_path, size=font_size) + except Exception: + font = ImageFont.load_default() + + # Draw bounding boxes + if detections: + colors = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#FF00FF', '#00FFFF', '#FFA500', '#800080'] + for idx, det in enumerate(detections): + bbox = det['bbox'] + label = det.get('label', 'object') + confidence = det.get('confidence', 1.0) + + # Convert normalized coordinates to pixel coordinates + x1 = int(bbox[0] * width) + y1 = int(bbox[1] * height) + x2 = int(bbox[2] * width) + y2 = int(bbox[3] * height) + + # Choose color + color = colors[idx % len(colors)] + + # Draw box + draw.rectangle([x1, y1, x2, y2], outline=color, width=max(2, int(min(width, height) * 0.003))) + + # Draw label with background + label_text = f"{label} {confidence:.2f}" if confidence < 1.0 else label + bbox_font = draw.textbbox((x1, y1), label_text, font=font) + text_width = bbox_font[2] - bbox_font[0] + text_height = bbox_font[3] - bbox_font[1] + draw.rectangle([x1, y1 - text_height - 4, x1 + text_width + 4, y1], fill=color) + draw.text((x1 + 2, y1 - text_height - 2), label_text, fill='white', font=font) + + # Draw points + if points: + point_radius = max(3, int(min(width, height) * 0.01)) + for px, py in points: + x = int(px * width) + y = int(py * height) + # Draw point as a circle + draw.ellipse([x - point_radius, y - point_radius, x + point_radius, y + point_radius], + fill='#FF0000', outline='#FFFFFF', width=2) + + return annotated diff --git a/modules/interrogate/vqa_draw.py b/modules/interrogate/vqa_draw.py deleted file mode 100644 index 6b8bbcf7e..000000000 --- a/modules/interrogate/vqa_draw.py +++ /dev/null @@ -1,76 +0,0 @@ -# VQA Image Annotation Utilities -# Drawing functions for bounding boxes, points, and other visual annotations - -from PIL import Image, ImageDraw, ImageFont -from modules import shared - - -def draw_bounding_boxes(image: Image.Image, detections: list, points: list = None) -> Image.Image: - """ - Draw bounding boxes and/or points on an image. - - Args: - image: PIL Image to annotate - detections: List of detection dicts with format: - [{'label': str, 'bbox': [x1, y1, x2, y2], 'confidence': float}, ...] - where coordinates are normalized 0-1 - points: Optional list of (x, y) tuples with normalized 0-1 coordinates - - Returns: - Annotated PIL Image with boxes and labels drawn, or None if no annotations - """ - if not detections and not points: - return None - - # Create a copy to avoid modifying original - annotated = image.copy() - draw = ImageDraw.Draw(annotated) - width, height = image.size - - # Try to load a font, fall back to default if unavailable - try: - font_size = max(12, int(min(width, height) * 0.02)) - font_path = shared.opts.font or "javascript/notosans-nerdfont-regular.ttf" - font = ImageFont.truetype(font_path, size=font_size) - except Exception: - font = ImageFont.load_default() - - # Draw bounding boxes - if detections: - colors = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#FF00FF', '#00FFFF', '#FFA500', '#800080'] - for idx, det in enumerate(detections): - bbox = det['bbox'] - label = det.get('label', 'object') - confidence = det.get('confidence', 1.0) - - # Convert normalized coordinates to pixel coordinates - x1 = int(bbox[0] * width) - y1 = int(bbox[1] * height) - x2 = int(bbox[2] * width) - y2 = int(bbox[3] * height) - - # Choose color - color = colors[idx % len(colors)] - - # Draw box - draw.rectangle([x1, y1, x2, y2], outline=color, width=max(2, int(min(width, height) * 0.003))) - - # Draw label with background - label_text = f"{label} {confidence:.2f}" if confidence < 1.0 else label - bbox_font = draw.textbbox((x1, y1), label_text, font=font) - text_width = bbox_font[2] - bbox_font[0] - text_height = bbox_font[3] - bbox_font[1] - draw.rectangle([x1, y1 - text_height - 4, x1 + text_width + 4, y1], fill=color) - draw.text((x1 + 2, y1 - text_height - 2), label_text, fill='white', font=font) - - # Draw points - if points: - point_radius = max(3, int(min(width, height) * 0.01)) - for px, py in points: - x = int(px * width) - y = int(py * height) - # Draw point as a circle - draw.ellipse([x - point_radius, y - point_radius, x + point_radius, y + point_radius], - fill='#FF0000', outline='#FFFFFF', width=2) - - return annotated