| |
| |
| |
| |
| |
|
|
| 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 |
|
|
| 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") |
|
|
| |
| |
| 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) |
|
|
