441 lines
15 KiB
Python
441 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Lightweight image-to-text API for Jetson Nano + EdgeTPU.
|
|
Compatible with Python 3.6.
|
|
"""
|
|
import argparse
|
|
import base64
|
|
import io
|
|
import os
|
|
import sys
|
|
import time
|
|
import uuid
|
|
import threading
|
|
from flask import Flask, request, jsonify
|
|
|
|
from core import image_io, coral_vision, text_detect, ocr, schema
|
|
|
|
try:
|
|
import cv2
|
|
except ImportError: # pragma: no cover
|
|
cv2 = None
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
|
|
# ---- Configuration ----
|
|
|
|
def env_bool(name, default):
|
|
val = os.getenv(name)
|
|
if val is None:
|
|
return default
|
|
return val.lower() in ("1", "true", "yes", "on")
|
|
|
|
def env_float(name, default):
|
|
val = os.getenv(name)
|
|
try:
|
|
return float(val) if val is not None else default
|
|
except ValueError:
|
|
return default
|
|
|
|
def env_int(name, default):
|
|
val = os.getenv(name)
|
|
try:
|
|
return int(val) if val is not None else default
|
|
except ValueError:
|
|
return default
|
|
|
|
MODEL_EDGETPU_PATH = os.getenv("MODEL_EDGETPU_PATH", "./models/ssd_mobilenet_v2_coco_quant_postprocess_edgetpu.tflite")
|
|
MODEL_LABELS_PATH = os.getenv("MODEL_LABELS_PATH", "./models/coco_labels.txt")
|
|
MODEL_EAST_PATH = os.getenv("MODEL_EAST_PATH", "./models/frozen_east_text_detection.pb")
|
|
OCR_LANG = os.getenv("OCR_LANG", "eng+rus")
|
|
MAX_IMAGE_MB = env_int("MAX_IMAGE_MB", 10)
|
|
CONF_THRESHOLD = env_float("CONF_THRESHOLD", 0.4)
|
|
RETURN_BOXES = env_bool("RETURN_BOXES", True)
|
|
OCR_CONFIG = os.getenv("OCR_CONFIG", "--oem 1 --psm 6")
|
|
TESSDATA_PREFIX = os.getenv("TESSDATA_PREFIX")
|
|
USE_CUDA_DNN = env_bool("USE_CUDA_DNN", False)
|
|
DEBUG_VISUAL = env_bool("DEBUG_VISUAL", False)
|
|
SHOW_GUI = env_bool("SHOW_GUI", False)
|
|
GUI_TARGET_WIDTH = env_int("GUI_TARGET_WIDTH", 1600)
|
|
OCR_UPSCALE = env_float("OCR_UPSCALE", 1.5) # 1.0 = no upscale
|
|
OCR_CLAHE = env_bool("OCR_CLAHE", True)
|
|
OCR_BILATERAL = env_bool("OCR_BILATERAL", False)
|
|
OCR_BLUR = env_int("OCR_BLUR", 0) # 0 disables gaussian blur
|
|
OCR_INVERT = env_bool("OCR_INVERT", False)
|
|
OCR_USE_THRESHOLD = env_bool("OCR_USE_THRESHOLD", True)
|
|
OCR_AUTO_LANG = env_bool("OCR_AUTO_LANG", True)
|
|
MAX_CONTENT_LENGTH = MAX_IMAGE_MB * 1024 * 1024
|
|
|
|
app = Flask(__name__)
|
|
app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH
|
|
|
|
if not TESSDATA_PREFIX:
|
|
# Attempt common system locations
|
|
for candidate in (
|
|
"/usr/share/tesseract-ocr/4.00/tessdata_best",
|
|
"/usr/share/tesseract-ocr/4.00/tessdata",
|
|
):
|
|
if os.path.exists(candidate):
|
|
TESSDATA_PREFIX = candidate
|
|
break
|
|
|
|
if TESSDATA_PREFIX:
|
|
os.environ["TESSDATA_PREFIX"] = TESSDATA_PREFIX
|
|
|
|
_GUI_FRAME = None
|
|
_GUI_LOCK = threading.Lock()
|
|
_GUI_THREAD = None
|
|
|
|
def _start_gui_thread():
|
|
if not SHOW_GUI or cv2 is None:
|
|
return
|
|
global _GUI_THREAD
|
|
if _GUI_THREAD is not None:
|
|
return
|
|
def _loop():
|
|
cv2.namedWindow("image2text debug", cv2.WINDOW_NORMAL)
|
|
while True:
|
|
frame = None
|
|
with _GUI_LOCK:
|
|
if _GUI_FRAME is not None:
|
|
frame = _GUI_FRAME.copy()
|
|
if frame is not None:
|
|
cv2.imshow("image2text debug", frame)
|
|
cv2.resizeWindow("image2text debug", frame.shape[1], frame.shape[0])
|
|
cv2.waitKey(50)
|
|
_GUI_THREAD = threading.Thread(target=_loop, daemon=True)
|
|
_GUI_THREAD.start()
|
|
|
|
def _show_placeholder():
|
|
if not SHOW_GUI or cv2 is None:
|
|
return
|
|
_gui_show_message("image2text waiting for request")
|
|
|
|
|
|
def _gui_show(frame):
|
|
if not SHOW_GUI or cv2 is None or frame is None:
|
|
return
|
|
global _GUI_FRAME
|
|
with _GUI_LOCK:
|
|
_GUI_FRAME = frame
|
|
|
|
|
|
def _gui_show_message(msg):
|
|
if not SHOW_GUI or cv2 is None:
|
|
return
|
|
width = GUI_TARGET_WIDTH if GUI_TARGET_WIDTH > 0 else 1280
|
|
height = int(width * 9 / 16)
|
|
canvas = np.ones((height, width, 3), dtype=np.uint8) * 240
|
|
cv2.putText(canvas, msg[:60], (40, height // 2),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2, cv2.LINE_AA)
|
|
_gui_show(canvas)
|
|
|
|
# ---- Load models once ----
|
|
|
|
coral = coral_vision.CoralVision(
|
|
model_path=MODEL_EDGETPU_PATH,
|
|
labels_path=MODEL_LABELS_PATH,
|
|
score_threshold=CONF_THRESHOLD,
|
|
)
|
|
|
|
txt_detector = text_detect.TextDetector(
|
|
east_model_path=MODEL_EAST_PATH,
|
|
score_threshold=CONF_THRESHOLD,
|
|
use_cuda=USE_CUDA_DNN,
|
|
)
|
|
|
|
# ---- Helpers ----
|
|
|
|
def parse_image_from_request(req):
|
|
if 'image' in req.files:
|
|
data = req.files['image'].read()
|
|
source = 'multipart'
|
|
else:
|
|
try:
|
|
payload = req.get_json(force=True, silent=False)
|
|
except Exception:
|
|
payload = None
|
|
if payload and 'image_base64' in payload:
|
|
b64_str = payload['image_base64'] or ''
|
|
if ',' in b64_str:
|
|
b64_str = b64_str.split(',', 1)[1]
|
|
try:
|
|
data = base64.b64decode(b64_str, validate=False)
|
|
except Exception:
|
|
return None, 'invalid_base64'
|
|
source = 'base64'
|
|
else:
|
|
return None, 'not_found'
|
|
if not data:
|
|
return None, 'empty'
|
|
if len(data) > MAX_CONTENT_LENGTH:
|
|
return None, 'too_large'
|
|
return data, source
|
|
|
|
|
|
def build_debug_image(img_rgb, boxes, blocks, objects, target_width=1024, return_image=False):
|
|
"""Return (base64 JPEG, optional ndarray for display) with drawn boxes."""
|
|
try:
|
|
vis = img_rgb.copy()
|
|
h, w = vis.shape[:2]
|
|
scale = 1.0
|
|
if target_width and target_width > 0 and w != target_width:
|
|
scale = float(target_width) / float(w)
|
|
# avoid crazy upscales
|
|
if scale > 3.0:
|
|
scale = 3.0
|
|
new_w = int(w * scale)
|
|
new_h = int(h * scale)
|
|
vis = np.array(Image.fromarray(vis).resize((new_w, new_h)))
|
|
# draw using cv2 if available
|
|
if cv2 is not None:
|
|
vis_bgr = cv2.cvtColor(vis, cv2.COLOR_RGB2BGR)
|
|
# text boxes in blue
|
|
for b in boxes:
|
|
x1, y1, x2, y2 = [int(v * scale) for v in b]
|
|
cv2.rectangle(vis_bgr, (x1, y1), (x2, y2), (255, 0, 0), 2)
|
|
for blk in blocks:
|
|
x1, y1, x2, y2 = [int(v * scale) for v in blk.get("bbox", [0, 0, 0, 0])]
|
|
cv2.rectangle(vis_bgr, (x1, y1), (x2, y2), (0, 128, 255), 2)
|
|
txt = blk.get("text", "")[:20]
|
|
cv2.putText(vis_bgr, txt, (x1, max(0, y1 - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 128, 255), 1, cv2.LINE_AA)
|
|
# objects in green
|
|
for obj in objects or []:
|
|
bbox = obj.get("bbox", [0, 0, 0, 0])
|
|
x1, y1, x2, y2 = [int(v * scale) for v in bbox]
|
|
cv2.rectangle(vis_bgr, (x1, y1), (x2, y2), (0, 200, 0), 2)
|
|
lbl = obj.get("label", "obj")
|
|
cv2.putText(vis_bgr, lbl, (x1, max(0, y1 - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 200, 0), 1, cv2.LINE_AA)
|
|
ok, buf = cv2.imencode(".jpg", vis_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
|
|
if not ok:
|
|
return None, None
|
|
b64 = base64.b64encode(buf.tobytes()).decode("ascii")
|
|
return b64, vis_bgr if return_image else None
|
|
# fallback to PIL
|
|
img = Image.fromarray(vis)
|
|
draw = ImageDraw.Draw(img)
|
|
for b in boxes:
|
|
x1, y1, x2, y2 = [int(v * scale) for v in b]
|
|
draw.rectangle([x1, y1, x2, y2], outline="blue", width=2)
|
|
for blk in blocks:
|
|
x1, y1, x2, y2 = [int(v * scale) for v in blk.get("bbox", [0, 0, 0, 0])]
|
|
draw.rectangle([x1, y1, x2, y2], outline="orange", width=2)
|
|
txt = blk.get("text", "")[:20]
|
|
draw.text((x1, max(0, y1 - 10)), txt, fill="orange")
|
|
for obj in objects or []:
|
|
bbox = obj.get("bbox", [0, 0, 0, 0])
|
|
x1, y1, x2, y2 = [int(v * scale) for v in bbox]
|
|
draw.rectangle([x1, y1, x2, y2], outline="green", width=2)
|
|
draw.text((x1, max(0, y1 - 10)), obj.get("label", "obj"), fill="green")
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="JPEG", quality=85)
|
|
b64 = base64.b64encode(buf.getvalue()).decode("ascii")
|
|
return b64, np.array(img) if return_image else None
|
|
except Exception:
|
|
return None, None
|
|
|
|
# ---- Routes ----
|
|
|
|
@app.route('/health', methods=['GET'])
|
|
def health():
|
|
return jsonify({"status": "ok", "coral": coral.available})
|
|
|
|
|
|
@app.route('/v1/image2text', methods=['POST'])
|
|
def image2text():
|
|
t0 = time.time()
|
|
request_id = str(uuid.uuid4())
|
|
data, status = parse_image_from_request(request)
|
|
if SHOW_GUI and cv2 is not None:
|
|
_gui_show_message("processing request {}".format(request_id[:8]))
|
|
|
|
if data is None:
|
|
if status == 'too_large':
|
|
if SHOW_GUI and cv2 is not None:
|
|
_gui_show_message("error: payload too large")
|
|
return jsonify({"error": "payload too large"}), 413
|
|
if status in ('not_found', 'empty'):
|
|
if SHOW_GUI and cv2 is not None:
|
|
_gui_show_message("error: image missing")
|
|
return jsonify({"error": "image is required"}), 400
|
|
if SHOW_GUI and cv2 is not None:
|
|
_gui_show_message("error: invalid base64")
|
|
return jsonify({"error": "invalid image_base64"}), 400
|
|
|
|
img_rgb = image_io.decode_image(data)
|
|
if img_rgb is None:
|
|
if SHOW_GUI and cv2 is not None:
|
|
_gui_show_message("error: unsupported format")
|
|
return jsonify({"error": "unsupported image format"}), 415
|
|
|
|
objects = []
|
|
coral_used = False
|
|
try:
|
|
objects, coral_used = coral.infer(img_rgb)
|
|
except Exception as exc:
|
|
print("[WARN] Coral failure {}: {}".format(request_id, exc), file=sys.stderr)
|
|
objects = []
|
|
coral_used = False
|
|
|
|
boxes, east_used = txt_detector.detect(img_rgb)
|
|
|
|
prep_options = {
|
|
"upscale": OCR_UPSCALE,
|
|
"clahe": OCR_CLAHE,
|
|
"bilateral": OCR_BILATERAL,
|
|
"blur": OCR_BLUR,
|
|
"invert": OCR_INVERT,
|
|
"use_threshold": OCR_USE_THRESHOLD,
|
|
}
|
|
|
|
ocr_lang_selected = OCR_LANG
|
|
ocr_lang_strategy = "fixed"
|
|
if OCR_AUTO_LANG:
|
|
try:
|
|
ocr_lang_selected, ocr_lang_strategy = ocr.choose_ocr_lang(
|
|
img_rgb,
|
|
boxes,
|
|
requested_lang=OCR_LANG,
|
|
prep_options=prep_options,
|
|
)
|
|
except Exception as exc:
|
|
print("[WARN] OCR lang auto-detect failed {}: {}".format(request_id, exc), file=sys.stderr)
|
|
ocr_lang_selected = OCR_LANG
|
|
ocr_lang_strategy = "auto_error_fallback"
|
|
|
|
text, blocks = ocr.run_ocr(
|
|
img_rgb,
|
|
boxes,
|
|
lang=ocr_lang_selected,
|
|
config=OCR_CONFIG,
|
|
return_boxes=RETURN_BOXES,
|
|
prep_options=prep_options,
|
|
)
|
|
ocr_passes = 1
|
|
|
|
# Retry with requested language set when auto-picked language produced no text.
|
|
if (not (text or '').strip()) and ocr_lang_selected != OCR_LANG:
|
|
retry_text, retry_blocks = ocr.run_ocr(
|
|
img_rgb,
|
|
boxes,
|
|
lang=OCR_LANG,
|
|
config=OCR_CONFIG,
|
|
return_boxes=RETURN_BOXES,
|
|
prep_options=prep_options,
|
|
)
|
|
ocr_passes += 1
|
|
if (retry_text or '').strip():
|
|
text, blocks = retry_text, retry_blocks
|
|
ocr_lang_selected = OCR_LANG
|
|
ocr_lang_strategy = "{}+lang_retry".format(ocr_lang_strategy)
|
|
|
|
# Retry with full-frame OCR if EAST boxes were too aggressive and OCR is empty.
|
|
if not (text or '').strip():
|
|
h, w = img_rgb.shape[:2]
|
|
full_box = (0, 0, w, h)
|
|
needs_full_pass = True
|
|
if len(boxes) == 1:
|
|
try:
|
|
cur_box = tuple([int(v) for v in boxes[0]])
|
|
needs_full_pass = cur_box != full_box
|
|
except Exception:
|
|
needs_full_pass = True
|
|
if needs_full_pass:
|
|
retry_text, retry_blocks = ocr.run_ocr(
|
|
img_rgb,
|
|
[full_box],
|
|
lang=ocr_lang_selected,
|
|
config=OCR_CONFIG,
|
|
return_boxes=RETURN_BOXES,
|
|
prep_options=prep_options,
|
|
)
|
|
ocr_passes += 1
|
|
if (retry_text or '').strip():
|
|
text, blocks = retry_text, retry_blocks
|
|
ocr_lang_strategy = "{}+full_frame_retry".format(ocr_lang_strategy)
|
|
|
|
scene_summary = schema.summarize_scene(objects)
|
|
|
|
latency_ms = int((time.time() - t0) * 1000)
|
|
|
|
debug_requested = DEBUG_VISUAL or request.args.get('debug', '').lower() in ('1', 'true', 'yes')
|
|
if SHOW_GUI:
|
|
debug_requested = True # ensure overlay built when GUI requested
|
|
debug_image_b64 = None
|
|
debug_image_arr = None
|
|
if debug_requested:
|
|
target_w = GUI_TARGET_WIDTH if SHOW_GUI else 1024
|
|
debug_image_b64, debug_image_arr = build_debug_image(
|
|
img_rgb,
|
|
boxes,
|
|
blocks if RETURN_BOXES else [],
|
|
objects,
|
|
target_width=target_w,
|
|
return_image=SHOW_GUI,
|
|
)
|
|
if SHOW_GUI and cv2 is not None and debug_image_arr is not None:
|
|
try:
|
|
_gui_show(debug_image_arr)
|
|
except Exception:
|
|
pass
|
|
|
|
resp = schema.build_response(
|
|
text=text,
|
|
blocks=blocks,
|
|
objects=objects,
|
|
scene_summary=scene_summary,
|
|
metadata={
|
|
"latency_ms": latency_ms,
|
|
"coral_used": coral_used,
|
|
"east_used": east_used,
|
|
"ocr_lang": ocr_lang_selected,
|
|
"ocr_lang_requested": OCR_LANG,
|
|
"ocr_lang_strategy": ocr_lang_strategy,
|
|
"ocr_passes": ocr_passes,
|
|
"request_id": request_id,
|
|
"debug": debug_requested,
|
|
},
|
|
return_boxes=RETURN_BOXES,
|
|
)
|
|
|
|
if debug_image_b64:
|
|
resp["debug_image_base64"] = debug_image_b64
|
|
|
|
print("[INFO] req={} size={}KB coral={} east={} latency={}ms".format(
|
|
request_id,
|
|
len(data) // 1024,
|
|
coral_used,
|
|
east_used,
|
|
latency_ms
|
|
))
|
|
|
|
return jsonify(resp)
|
|
|
|
|
|
# ---- Entrypoint ----
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Image to text API')
|
|
parser.add_argument('--host', default='0.0.0.0')
|
|
parser.add_argument('--port', type=int, default=8000)
|
|
parser.add_argument('--debug-visual', action='store_true', help='force returning debug overlay base64')
|
|
parser.add_argument('--show-gui', action='store_true', help='show overlay in a window (requires DISPLAY and OpenCV GUI); forces debug overlay creation')
|
|
args = parser.parse_args()
|
|
|
|
global DEBUG_VISUAL
|
|
if args.debug_visual:
|
|
DEBUG_VISUAL = True
|
|
global SHOW_GUI
|
|
if args.show_gui:
|
|
SHOW_GUI = True
|
|
if SHOW_GUI:
|
|
_start_gui_thread()
|
|
_show_placeholder()
|
|
|
|
app.run(host=args.host, port=args.port, threaded=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|