diff --git a/tools/ora_editor/app.py b/tools/ora_editor/app.py index 697f2dd..99ddd32 100644 --- a/tools/ora_editor/app.py +++ b/tools/ora_editor/app.py @@ -11,15 +11,18 @@ from pathlib import Path from xml.etree import ElementTree as ET try: - from flask import ( - Flask, request, jsonify, send_file, Response, render_template, - make_response - ) -except ImportError: - print("Error: Flask is required. Install with: pip install -r requirements.txt") - sys.exit(1) +from flask import ( + Flask, request, jsonify, send_file, Response, render_template, + make_response +) from PIL import Image, ImageDraw +import logging +import urllib.parse + +# Configure logging +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) # Configure paths APP_DIR = Path(__file__).parent @@ -309,16 +312,23 @@ def api_polygon(): for field in required: if field not in data: return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400 - - _polygon_storage[data['ora_path']] = { - 'points': data['points'], - 'color': data.get('color', '#FF0000'), - 'width': data.get('width', 2) + + ora_path = data['ora_path'] + points = data['points'] + color = data.get('color', '#FF0000') + width = data.get('width', 2) + + logger.info(f"[POLYGON] Storing polygon: {len(points)} points, color: {color}, width: {width}") + + _polygon_storage[ora_path] = { + 'points': points, + 'color': color, + 'width': width } return jsonify({ 'success': True, - 'overlay_url': f'/api/image/polygon?ora_path={data["ora_path"]}' + 'overlay_url': f'/api/image/polygon?ora_path={ora_path}' }) @@ -333,112 +343,183 @@ def api_polygon_clear(): return jsonify({'success': True}) -@app.route('/api/mask/extract', methods=['POST']) + @app.route('/api/mask/extract', methods=['POST']) def api_mask_extract(): """Extract mask using ComfyUI.""" data = request.get_json() - required = ['subject'] + required = ['subject', 'ora_path'] for field in required: if field not in data: return jsonify({'success': False, 'error': f'Missing {field} parameter'}), 400 - + subject = data['subject'] use_polygon = data.get('use_polygon', False) - ora_path = data.get('ora_path') + ora_path = data['ora_path'] comfy_url = data.get('comfy_url', '127.0.0.1:8188') - - # Build the ComfyUI prompt - import json - import urllib.request - import base64 - import uuid as uuid_lib - + + logger.info(f"[MASK EXTRACT] Subject: {subject}") + logger.info(f"[MASK EXTRACT] Use polygon: {use_polygon}") + logger.info(f"[MASK EXTRACT] ORA path: {ora_path}") + logger.info(f"[MASK EXTRACT] ComfyUI URL: {comfy_url}") + + # Check if we need to apply polygon to image + if use_polygon: + polygon_data = _polygon_storage.get(ora_path, {}) + polygon_points = polygon_data.get('points', []) + + if not polygon_points or len(polygon_points) < 3: + logger.warning(f"[MASK EXTRACT] Use polygon requested but no valid polygon points stored") + return jsonify({'success': False, 'error': 'No valid polygon points. Please draw polygon first.'}), 400 + + logger.info(f"[MASK EXTRACT] Applying polygon overlay: {len(polygon_points)} points") + # Load workflow template workflow_path = APP_DIR.parent / "image_mask_extraction.json" if not workflow_path.exists(): + logger.error(f"[MASK EXTRACT] Workflow file not found: {workflow_path}") return jsonify({'success': False, 'error': f'Workflow file not found: {workflow_path}'}), 500 - + with open(workflow_path) as f: workflow = json.load(f) - + + logger.info(f"[MASK EXTRACT] Loaded workflow") + # Update prompt text in workflow + + # Load base image from ORA + base_img = None + try: + with zipfile.ZipFile(ora_path, 'r') as zf: + img_data = zf.read('mergedimage.png') + base_img = Image.open(io.BytesIO(img_data)).convert('RGBA') + logger.info(f"[MASK EXTRACT] Loaded base image: {base_img.size}") + + except Exception as e: + logger.error(f"[MASK EXTRACT] Error loading base image: {e}") + return jsonify({'success': False, 'error': f'Error loading image: {str(e)}'}), 500 + + # Apply polygon overlay if requested + if use_polygon: + polygon_data = _polygon_storage.get(ora_path, {}) + polygon_points = polygon_data.get('points', []) + + if not polygon_points or len(polygon_points) < 3: + logger.warning(f"[MASK EXTRACT] Use polygon requested but no valid polygon points stored") + return jsonify({'success': False, 'error': 'No valid polygon points. Please draw polygon first.'}), 400 + + logger.info(f"[MASK EXTRACT] Applying polygon overlay: {len(polygon_points)} points") + + # Convert percentage points to pixel coordinates + w, h = base_img.size + pixel_points = [(int(p['x'] * w), int(p['y'] * h)) for p in polygon_points] + + # Draw polygon on image + polygon_color = polygon_data.get('color', '#FF0000') + polygon_width = polygon_data.get('width', 2) + + draw = ImageDraw.Draw(base_img) + + # Draw the polygon + hex_color = polygon_color if len(polygon_color) == 7 else polygon_color + 'FF' + draw.polygon(pixel_points, outline=hex_color, width=polygon_width) + + logger.info(f"[MASK EXTRACT] Drew polygon on image") + + # Encode (possibly modified) image + img_io = io.BytesIO() + base_img.save(img_io, format='PNG') + img_io.seek(0) + + base64_image = base64.b64encode(img_io.read()).decode('utf-8') + logger.info(f"[MASK EXTRACT] Encoded image as base64: {len(base64_image)} chars") + + # Set image input in workflow + workflow["168"]["inputs"]["image"] = base64_image + + # Update prompt text for node_id, node in workflow.items(): if 'inputs' in node and 'text' in node['inputs']: node['inputs']['text'] = f"Create a black and white alpha mask of {subject}, leaving everything else black" break - - # Queue the prompt - prompt_id = str(uuid_lib.uuid4()) - headers = {'Content-Type': 'application/json'} - - try: - req_data = json.dumps({'prompt': workflow, 'client_id': prompt_id}).encode() - req = urllib.request.Request( - f'http://{comfy_url}/prompt', - data=req_data, - headers=headers, - method='POST' - ) - - with urllib.request.urlopen(req, timeout=30) as response: - result = json.loads(response.read().decode()) - - except Exception as e: - return jsonify({'success': False, 'error': f'Failed to connect to ComfyUI: {str(e)}'}), 500 - - # Poll for completion (up to 4 minutes) - start_time = time.time() - timeout = 240 # 4 minutes - - while time.time() - start_time < timeout: - try: - req = urllib.request.Request( - f'http://{comfy_url}/history/{prompt_id}', - headers=headers, - method='GET' - ) - - with urllib.request.urlopen(req, timeout=30) as response: - history = json.loads(response.read().decode()) - if prompt_id in history: - outputs = history[prompt_id]['outputs'] - for node_id, node_output in outputs.items(): - images = node_output.get('images', []) - for img_info in images: - filename = img_info['filename'] - subfolder = img_info.get('subfolder', '') + # Poll for completion (up to 4 minutes) + start_time = time.time() + timeout = 240 + + while time.time() - start_time < timeout: + try: + req = urllib.request.Request( + f'http://{comfy_url}/history/{prompt_id}', + headers=headers, + method='GET' + ) + + with urllib.request.urlopen(req, timeout=30) as response: + history = json.loads(response.read().decode()) - # Download the image - download_req = urllib.request.Request( - f'http://{comfy_url}/view?filename={filename}&subfolder={subfolder}&type=output', - headers=headers, - method='GET' - ) - - with urllib.request.urlopen(download_req, timeout=30) as img_response: - img_data = img_response.read() - - # Save to temp directory - timestamp = str(int(time.time())) - mask_path = TEMP_DIR / f"mask_{timestamp}.png" - with open(mask_path, 'wb') as f: - f.write(img_data) - - return jsonify({ - 'success': True, - 'mask_path': str(mask_path), - 'mask_url': f'/api/file/mask?path={mask_path}' - }) - - except urllib.error.HTTPError as e: - if e.code != 404: - break + if prompt_id in history: + outputs = history[prompt_id].get('outputs', {}) + + logger.info(f"[MASK EXTRACT] Status check - prompt_id found in history") + + for node_id, node_output in outputs.items(): + if 'images' in node_output: + images = node_output['images'] + + logger.info(f"[MASK EXTRACT] Found {len(images)} images in node {node_id}") + + for img_info in images: + filename = img_info.get('filename', '') + subfolder = img_info.get('subfolder', '') + + logger.info(f"[MASK EXTRACT] Downloading image: {filename} from subfolder: {subfolder}") + + # Download the image + params = { + 'filename': filename, + 'subfolder': subfolder, + 'type': 'output' + } + + download_req = urllib.request.Request( + f'http://{comfy_url}/view?{urllib.parse.urlencode(params)}', + headers=headers, + method='GET' + ) + + with urllib.request.urlopen(download_req, timeout=30) as img_response: + img_data = img_response.read() + + # Save to temp directory + timestamp = str(int(time.time())) + mask_path = TEMP_DIR / f"mask_{timestamp}.png" + + with open(mask_path, 'wb') as f: + f.write(img_data) + + logger.info(f"[MASK EXTRACT] Saved mask to: {mask_path} ({len(img_data)} bytes)") + + return jsonify({ + 'success': True, + 'mask_path': str(mask_path), + 'mask_url': f'/api/file/mask?path={mask_path}' + }) + else: + logger.info(f"[MASK EXTRACT] Prompt_id not in history yet, waiting...") + except urllib.error.HTTPError as e: + if e.code != 404: + logger.error(f"[MASK EXTRACT] HTTP Error polling: {e.code}") + break + time.sleep(2) + + # Timeout + logger.error(f"[MASK EXTRACT] Timeout waiting for ComfyUI to complete") + return jsonify({'success': False, 'error': 'Mask extraction timed out'}), 500 - time.sleep(2) - - return jsonify({'success': False, 'error': 'Mask extraction timed out'}), 500 + except Exception as e: + logger.exception(f"[MASK EXTRACT] Error: {e}") + return jsonify({'success': False, 'error': str(e)}), 500 @app.route('/api/krita/open', methods=['POST'])