Text Recognition (OCR)

1import os2import sys3import json4import traceback56from gais import Gais78INPUT_DIR = "/data/input"91011def main():12    try:13        input_json = sys.stdin.read()14        execution_input = json.loads(input_json)15        inputs = execution_input.get("inputs", {})1617        media = inputs.get("media", "")18        lang = inputs.get("lang", "en")19        text_det_thresh = float(inputs.get("text_det_thresh", 0.3))20        text_rec_score_thresh = float(inputs.get("text_rec_score_thresh", 0.5))21        video_sample_fps = float(inputs.get("video_sample_fps", 2.0))2223        if not media:24            raise ValueError("Media input is required")2526        media_path = os.path.join(INPUT_DIR, media)27        if not os.path.exists(media_path):28            raise FileNotFoundError(f"Input file not found: {media_path}")2930        print(31            f"Running OCR: lang={lang}, det_thresh={text_det_thresh}, "32            f"rec_thresh={text_rec_score_thresh}, video_fps={video_sample_fps}",33            file=sys.stderr,34        )3536        result = Gais.detect.text(37            media=media_path,38            lang=lang,39            text_det_thresh=text_det_thresh,40            text_rec_score_thresh=text_rec_score_thresh,41            video_sample_fps=video_sample_fps,42        )4344        result_data = json.loads(result.content)4546        frames = result_data.get("frames", [])47        total_text_regions = result_data.get("total_text_regions", 0)48        total_frames = result_data.get("total_frames", 1)49        model_used = result_data.get("model_used", "PP-OCRv5")50        result_lang = result_data.get("lang", lang)51        input_type = result_data.get("input_type", "image")52        processing_time = result_data.get("processing_time_ms", 0)5354        print(55            f"OCR complete: {total_text_regions} text regions in {total_frames} frame(s), "56            f"model={model_used}, lang={result_lang}, type={input_type}, "57            f"processing_time={processing_time}ms",58            file=sys.stderr,59        )6061        # Flat output — keys match OUTPUT_SCHEMA62        output = {63            "ocr_results": json.dumps(frames),64            "total_text_regions": total_text_regions,65            "total_frames": total_frames,66            "model_used": model_used,67            "lang": result_lang,68            "input_type": input_type,69        }70        print(json.dumps(output, indent=2))7172    except Exception as e:73        error_output = {74            "error": str(e),75            "errorType": type(e).__name__,76            "traceback": traceback.format_exc(),77        }78        print(json.dumps(error_output), file=sys.stderr)79        sys.exit(1)808182if __name__ == "__main__":83    main()