feat(ai): add emit_progress() calls to all Python AI scripts

This commit is contained in:
Siddharth Kumar Sah
2026-03-23 01:38:19 +08:00
parent 7d74ddd3a6
commit 723842988b
5 changed files with 80 additions and 15 deletions
+13 -1
View File
@@ -3,6 +3,11 @@ import sys
import json import json
def emit_progress(percent, stage):
"""Emit structured progress to stderr for bridge.ts to capture."""
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
def main(): def main():
input_path = sys.argv[1] input_path = sys.argv[1]
output_path = sys.argv[2] output_path = sys.argv[2]
@@ -12,6 +17,7 @@ def main():
sensitivity = settings.get("sensitivity", 0.5) sensitivity = settings.get("sensitivity", 0.5)
try: try:
emit_progress(10, "Loading face detection model")
from PIL import Image, ImageFilter from PIL import Image, ImageFilter
img = Image.open(input_path).convert("RGB") img = Image.open(input_path).convert("RGB")
@@ -20,17 +26,21 @@ def main():
import mediapipe as mp import mediapipe as mp
import numpy as np import numpy as np
emit_progress(20, "Model ready")
mp_face = mp.solutions.face_detection mp_face = mp.solutions.face_detection
with mp_face.FaceDetection( with mp_face.FaceDetection(
min_detection_confidence=sensitivity min_detection_confidence=sensitivity
) as detector: ) as detector:
img_array = np.array(img) img_array = np.array(img)
emit_progress(25, "Scanning for faces")
results = detector.process(img_array) results = detector.process(img_array)
faces = [] faces = []
emit_progress(50, f"Found {len(results.detections or [])} faces")
if results.detections: if results.detections:
for detection in results.detections: for i, detection in enumerate(results.detections):
bbox = detection.location_data.relative_bounding_box bbox = detection.location_data.relative_bounding_box
x = int(bbox.xmin * img.width) x = int(bbox.xmin * img.width)
y = int(bbox.ymin * img.height) y = int(bbox.ymin * img.height)
@@ -50,7 +60,9 @@ def main():
) )
img.paste(blurred, (x1, y1)) img.paste(blurred, (x1, y1))
faces.append({"x": x, "y": y, "w": w, "h": h}) faces.append({"x": x, "y": y, "w": w, "h": h})
emit_progress(50 + int((i + 1) / max(len(results.detections), 1) * 40), f"Blurring face {i + 1} of {len(results.detections)}")
emit_progress(95, "Saving result")
img.save(output_path) img.save(output_path)
print( print(
json.dumps( json.dumps(
+12
View File
@@ -3,12 +3,18 @@ import sys
import json import json
def emit_progress(percent, stage):
"""Emit structured progress to stderr for bridge.ts to capture."""
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
def main(): def main():
input_path = sys.argv[1] input_path = sys.argv[1]
mask_path = sys.argv[2] mask_path = sys.argv[2]
output_path = sys.argv[3] output_path = sys.argv[3]
try: try:
emit_progress(10, "Loading inpainting model")
from PIL import Image from PIL import Image
try: try:
@@ -16,10 +22,13 @@ def main():
from lama_cleaner.model_manager import ModelManager from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import Config from lama_cleaner.schema import Config
emit_progress(20, "Model loaded")
img = Image.open(input_path).convert("RGB") img = Image.open(input_path).convert("RGB")
mask = Image.open(mask_path).convert("L") mask = Image.open(mask_path).convert("L")
# Resize mask to match image if needed # Resize mask to match image if needed
emit_progress(25, "Analyzing mask")
if mask.size != img.size: if mask.size != img.size:
mask = mask.resize(img.size, Image.NEAREST) mask = mask.resize(img.size, Image.NEAREST)
@@ -37,7 +46,10 @@ def main():
hd_strategy_crop_trigger_size=800, hd_strategy_crop_trigger_size=800,
hd_strategy_resize_limit=800, hd_strategy_resize_limit=800,
) )
emit_progress(40, "Inpainting region")
result = model_manager(img_array, mask_array, config) result = model_manager(img_array, mask_array, config)
emit_progress(85, "Refining edges")
emit_progress(95, "Saving result")
Image.fromarray(result).save(output_path) Image.fromarray(result).save(output_path)
method = "lama" method = "lama"
+12
View File
@@ -3,6 +3,11 @@ import sys
import json import json
def emit_progress(percent, stage):
"""Emit structured progress to stderr for bridge.ts to capture."""
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
def main(): def main():
input_path = sys.argv[1] input_path = sys.argv[1]
settings = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {} settings = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {}
@@ -11,12 +16,15 @@ def main():
language = settings.get("language", "en") language = settings.get("language", "en")
try: try:
emit_progress(10, "Loading OCR engine")
if engine == "paddleocr": if engine == "paddleocr":
try: try:
from paddleocr import PaddleOCR from paddleocr import PaddleOCR
ocr = PaddleOCR(use_angle_cls=True, lang=language) ocr = PaddleOCR(use_angle_cls=True, lang=language)
emit_progress(30, "Analyzing text regions")
result = ocr.ocr(input_path, cls=True) result = ocr.ocr(input_path, cls=True)
emit_progress(70, "Extracting text")
text = "\n".join( text = "\n".join(
[ [
line[1][0] line[1][0]
@@ -26,6 +34,7 @@ def main():
if line and line[1] if line and line[1]
] ]
) )
emit_progress(95, "Formatting results")
print( print(
json.dumps({"success": True, "text": text, "engine": "paddleocr"}) json.dumps({"success": True, "text": text, "engine": "paddleocr"})
) )
@@ -47,15 +56,18 @@ def main():
tess_lang = lang_map.get(language, "eng") tess_lang = lang_map.get(language, "eng")
try: try:
emit_progress(30, "Running Tesseract")
result = subprocess.run( result = subprocess.run(
["tesseract", input_path, "stdout", "-l", tess_lang], ["tesseract", input_path, "stdout", "-l", tess_lang],
capture_output=True, capture_output=True,
text=True, text=True,
timeout=120, timeout=120,
) )
emit_progress(70, "Extracting text")
text = result.stdout.strip() text = result.stdout.strip()
if result.returncode != 0 and not text: if result.returncode != 0 and not text:
raise RuntimeError(result.stderr.strip() or "Tesseract failed") raise RuntimeError(result.stderr.strip() or "Tesseract failed")
emit_progress(95, "Formatting results")
print( print(
json.dumps( json.dumps(
{"success": True, "text": text, "engine": "tesseract"} {"success": True, "text": text, "engine": "tesseract"}
+27 -10
View File
@@ -1,6 +1,12 @@
"""Background removal using rembg with state-of-the-art BiRefNet models.""" """Background removal using rembg with state-of-the-art BiRefNet models."""
import sys import sys
import json import json
import os
def emit_progress(percent, stage):
"""Emit structured progress to stderr for bridge.ts to capture."""
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
def main(): def main():
@@ -11,23 +17,26 @@ def main():
model = settings.get("model", "u2net") model = settings.get("model", "u2net")
bg_color = settings.get("backgroundColor", "") bg_color = settings.get("backgroundColor", "")
# Redirect stdout to stderr so library download/progress output
# cannot contaminate our JSON result on stdout.
stdout_fd = os.dup(1)
os.dup2(2, 1)
try: try:
from rembg import remove, new_session from rembg import remove, new_session
import io import io
# Progress messages go to stderr (stdout reserved for JSON result) emit_progress(10, "Loading model")
sys.stderr.write(f"Loading model: {model}\n")
sys.stderr.flush()
session = new_session(model) session = new_session(model)
sys.stderr.write("Processing image...\n") emit_progress(25, "Model loaded")
sys.stderr.flush()
with open(input_path, "rb") as f: with open(input_path, "rb") as f:
input_data = f.read() input_data = f.read()
# Try with alpha matting for better edges, fall back without # Try with alpha matting for better edges, fall back without
emit_progress(30, "Analyzing image")
try: try:
output_data = remove( output_data = remove(
input_data, input_data,
@@ -39,8 +48,11 @@ def main():
except Exception: except Exception:
output_data = remove(input_data, session=session) output_data = remove(input_data, session=session)
emit_progress(80, "Background removed")
# If a background color is specified, composite onto it # If a background color is specified, composite onto it
if bg_color and bg_color.startswith("#"): if bg_color and bg_color.startswith("#"):
emit_progress(85, "Compositing background")
from PIL import Image from PIL import Image
img = Image.open(io.BytesIO(output_data)).convert("RGBA") img = Image.open(io.BytesIO(output_data)).convert("RGBA")
@@ -54,22 +66,27 @@ def main():
bg.save(buf, format="PNG") bg.save(buf, format="PNG")
output_data = buf.getvalue() output_data = buf.getvalue()
emit_progress(95, "Saving result")
with open(output_path, "wb") as f: with open(output_path, "wb") as f:
f.write(output_data) f.write(output_data)
print(json.dumps({"success": True, "model": model})) result = json.dumps({"success": True, "model": model})
except ImportError: except ImportError:
print( result = json.dumps(
json.dumps(
{ {
"success": False, "success": False,
"error": "rembg is not installed. Install with: pip install rembg[cpu]", "error": "rembg is not installed. Install with: pip install rembg[cpu]",
} }
) )
)
except Exception as e: except Exception as e:
print(json.dumps({"success": False, "error": str(e)})) result = json.dumps({"success": False, "error": str(e)})
# Restore original stdout and write only our JSON result
os.dup2(stdout_fd, 1)
os.close(stdout_fd)
sys.stdout.write(result + "\n")
sys.stdout.flush()
if __name__ == "__main__": if __name__ == "__main__":
+12
View File
@@ -3,6 +3,11 @@ import sys
import json import json
def emit_progress(percent, stage):
"""Emit structured progress to stderr for bridge.ts to capture."""
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
def main(): def main():
input_path = sys.argv[1] input_path = sys.argv[1]
output_path = sys.argv[2] output_path = sys.argv[2]
@@ -11,6 +16,7 @@ def main():
scale = settings.get("scale", 2) scale = settings.get("scale", 2)
try: try:
emit_progress(10, "Loading upscale model")
from PIL import Image from PIL import Image
img = Image.open(input_path) img = Image.open(input_path)
@@ -36,14 +42,20 @@ def main():
model=model, model=model,
half=False, half=False,
) )
emit_progress(20, "Model ready")
img_array = np.array(img.convert("RGB")) img_array = np.array(img.convert("RGB"))
emit_progress(25, "Upscaling image")
output, _ = upsampler.enhance(img_array, outscale=scale) output, _ = upsampler.enhance(img_array, outscale=scale)
emit_progress(90, "Upscaling complete")
result = Image.fromarray(output) result = Image.fromarray(output)
emit_progress(95, "Saving result")
result.save(output_path) result.save(output_path)
method = "realesrgan" method = "realesrgan"
except (ImportError, Exception): except (ImportError, Exception):
# Fallback to Lanczos upscaling # Fallback to Lanczos upscaling
emit_progress(50, "Upscaling with Lanczos")
img_upscaled = img.resize(new_size, Image.LANCZOS) img_upscaled = img.resize(new_size, Image.LANCZOS)
emit_progress(95, "Saving result")
img_upscaled.save(output_path) img_upscaled.save(output_path)
method = "lanczos" method = "lanczos"