Add logging and improve mask extraction
- Added logging throughout the application - Updated mask extraction to load and pass base image to ComfyUI - Added polygon overlay support for mask extraction - Added urllib.parse import for URL encoding - Better error handling and status reporting
This commit is contained in:
@@ -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'])
|
||||
|
||||
Reference in New Issue
Block a user