# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# OCR for Text: DLStreamer pipeline sample.
# See README.md for prerequisites and inputs. Run: python3 dlstreamer_sample.py

import gi

gi.require_version("Gst", "1.0")
from gi.repository import Gst

import numpy as np
import openvino as ov

Gst.init([])

# Import cv2 after Gst.init to avoid GStreamer re-initialization conflicts.
import cv2

INPUT_VIDEO = "test_video.mp4"
DET_MODEL = "ch_PP-OCRv4_det_infer/inference.xml"
REC_MODEL = "ch_PP-OCRv4_rec_server_infer/inference.xml"
DICT_FILE = "ppocr_keys_v1.txt"
DET_SIZE = 960

core = ov.Core()
det_compiled = core.compile_model(core.read_model(DET_MODEL), "CPU")
rec_compiled = core.compile_model(core.read_model(REC_MODEL), "CPU")

# CTC label map: index 0 is the blank symbol, followed by every character in
# the dictionary file, followed by a trailing space character.
chars = open(DICT_FILE, encoding="utf-8").read().splitlines()
dict_character = ["blank"] + chars + [" "]


def detect_text_regions(frame, thresh=0.3, pad=4):
    """Return (x, y, w, h) boxes for words/lines of text in a frame.

    Resizing preserves aspect ratio (letterboxed onto a square canvas) so
    text isn't skewed, and dilating the detection map merges nearby
    characters into one box per word instead of one per character.
    """
    h0, w0 = frame.shape[:2]
    scale = DET_SIZE / max(h0, w0)
    resized = cv2.resize(frame, (int(w0 * scale), int(h0 * scale)))
    canvas = np.zeros((DET_SIZE, DET_SIZE, 3), dtype=np.uint8)
    canvas[:resized.shape[0], :resized.shape[1]] = resized

    blob = canvas.astype(np.float32).transpose(2, 0, 1)[np.newaxis] / 255.0
    det_map = det_compiled([blob])[det_compiled.output(0)][0, 0]
    binary = (det_map > thresh).astype(np.uint8) * 255
    dilated = cv2.dilate(binary, np.ones((9, 25), np.uint8))
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    boxes = []
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        if w < 10 or h < 5:
            continue
        x0, y0 = max(0, x / scale - pad), max(0, y / scale - pad)
        x1, y1 = min(w0, (x + w) / scale + pad), min(h0, (y + h) / scale + pad)
        boxes.append((int(x0), int(y0), int(x1 - x0), int(y1 - y0)))
    return boxes


def recognize_text(crop, rec_h=48, max_w=320):
    """Resize a cropped text region to the recognizer's input shape and
    CTC-decode the predicted character sequence into a string."""
    h, w = crop.shape[:2]
    if h == 0 or w == 0:
        return "", 0.0
    resized_w = max(1, min(max_w, round(rec_h * w / h)))
    blob = cv2.resize(crop, (resized_w, rec_h)).astype(np.float32) / 255.0
    blob = ((blob - 0.5) / 0.5).transpose(2, 0, 1)[np.newaxis, ...]

    preds = rec_compiled([blob])[rec_compiled.output(0)][0]
    idx = np.argmax(preds, axis=1)
    conf = np.max(preds, axis=1)

    text, scores, prev = [], [], -1
    for i, c in zip(idx, conf):
        if i != 0 and i != prev:
            text.append(dict_character[i])
            scores.append(c)
        prev = i
    confidence = float(np.mean(scores)) if scores else 0.0
    return "".join(text), confidence


def annotate(frame, box, text, confidence):
    """Draw a bounding box and, if any text was recognized, a legible
    label (solid background so it stays readable over any color) above it."""
    x, y, w, h = box
    cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
    if not text:
        return
    label = f"{text} ({confidence:.2f})"
    (tw, th), base = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
    top = max(0, y - th - base - 6)
    cv2.rectangle(frame, (x, top), (x + tw + 6, top + th + base + 6), (0, 255, 0), -1)
    cv2.putText(frame, label, (x + 3, top + th + 2),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)


pipeline_str = (
    f"filesrc location={INPUT_VIDEO} ! decodebin3 ! videoconvert ! "
    "video/x-raw,format=BGR ! "
    "appsink name=sink emit-signals=false sync=false"
)
pipeline = Gst.parse_launch(pipeline_str)
sink = pipeline.get_by_name("sink")
pipeline.set_state(Gst.State.PLAYING)

writer = None
frame_idx = 0
total_regions = 0

while True:
    sample = sink.emit("pull-sample")
    if sample is None:
        break
    buf = sample.get_buffer()
    caps = sample.get_caps().get_structure(0)
    width = caps.get_value("width")
    height = caps.get_value("height")

    ok, mapinfo = buf.map(Gst.MapFlags.READ)
    if not ok:
        continue
    frame = np.ndarray((height, width, 3), dtype=np.uint8,
                       buffer=mapinfo.data).copy()
    buf.unmap(mapinfo)
    frame_idx += 1

    for box in detect_text_regions(frame):
        x, y, w, h = box
        text, confidence = recognize_text(frame[y:y + h, x:x + w])
        annotate(frame, box, text, confidence)
        total_regions += 1
        print(f"Frame {frame_idx}: region=({x},{y},{w},{h}) text={text!r} "
              f"confidence={confidence:.2f}", flush=True)

    if writer is None:
        writer = cv2.VideoWriter(
            "output_dlstreamer.mp4", cv2.VideoWriter_fourcc(*"mp4v"),
            30.0, (width, height))
    writer.write(frame)

pipeline.set_state(Gst.State.NULL)
if writer:
    writer.release()
print(f"Total text regions across all frames: {total_regions}", flush=True)