mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
feat(ai): add emit_progress() calls to all Python AI scripts
This commit is contained in:
@@ -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(
|
||||||
|
|||||||
@@ -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"
|
||||||
|
|
||||||
|
|||||||
@@ -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"}
|
||||||
|
|||||||
@@ -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__":
|
||||||
|
|||||||
@@ -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"
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user