chore: wip

.gitattributes

@@ -0,0 +1 @@
+*.onnx filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+/.gradio/
+/.pytest_cache/
+/.venv/
+__pycache__

Dockerfile

@@ -0,0 +1,14 @@
+FROM python:3-bookworm
+
+RUN apt-get update && \
+    apt-get install -y libmagickwand-dev ghostscript fonts-freefont-ttf curl libgl1-mesa-glx
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+WORKDIR /app
+COPY --chown=user . /app
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    uv sync
+ENV PATH="/app/.venv/bin:$PATH"
+
+CMD ["python", "run.py"]

README.md

@@ -0,0 +1,8 @@
+---
+title: Image translator
+emoji: 🐳
+colorFrom: purple
+colorTo: gray
+sdk: docker
+app_port: 7860
+---

app/__init__.py

@@ -0,0 +1 @@
+from .utils import ocr

app/constants.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+
+FONT_PATH = "/usr/share/fonts/truetype/freefont/FreeMono.ttf"
+MODEL_PATH = str(Path(__file__).parent.parent / "model/en_infer.onnx")

app/utils.py

@@ -0,0 +1,32 @@
+from .constants import MODEL_PATH
+from rapidocr_onnxruntime import RapidOCR
+from numpy.typing import NDArray
+from os import PathLike
+from shapely import Polygon
+from shapely.geometry import box, MultiPolygon
+from shapely.ops import unary_union
+from typing import cast
+
+_rapid_ocr = RapidOCR(rec_model_path=MODEL_PATH)
+
+
+def ocr(image: NDArray | PathLike, esp: float = 0.5) -> list[tuple[Polygon, str]]:
+    result, _ = _rapid_ocr(image)
+    result.sort(key=lambda x: x[0][0][1])
+    texts: list[str] = []
+    shapes: list[Polygon] = []
+    for points, text, _ in result:
+        x, y, xmax, ymax = Polygon(points).bounds
+        ymax += (ymax - y) * esp
+        shapes.append(box(x, y, xmax, ymax))
+        texts.append(text)
+    merged_shapes = cast(MultiPolygon, unary_union(shapes)).geoms
+    merged_texts = [""] * len(merged_shapes)
+    visited = [False] * len(texts)
+    for i, shape in enumerate(merged_shapes):
+        for j, text in enumerate(texts):
+            if not visited[j]:
+                if shapes[j].intersects(shape):
+                    merged_texts[i] += f" {text}"
+                    visited[j] = True
+    return list(zip(merged_shapes, merged_texts))

model/en_infer.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef7abd8bd3629ae57ea2c28b425c1bd258a871b93fd2fe7c433946ade9b5d9ea
+size 8967018

pyproject.toml

@@ -0,0 +1,19 @@
+[project]
+name = "imgtran"
+version = "0.0.1"
+dependencies = [
+    "deep-translator>=1.11.4",
+    "fastapi[standard]>=0.115.6",
+    "gradio>=5.12.0",
+    "rapidocr-onnxruntime>=1.4.3",
+    "shapely>=2.0.6",
+    "wand>=0.6.13",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.3.4",
+]
+
+[tool.pytest.ini_options]
+pythonpath = '.'

run.py

@@ -0,0 +1,51 @@
+import cv2
+import io
+import gradio as gr
+from app import ocr
+from app.constants import FONT_PATH
+from numpy.typing import NDArray
+from deep_translator import GoogleTranslator
+from wand.color import Color
+from wand.drawing import Drawing
+from wand.image import Image
+from wand.font import Font
+from PIL import Image as PILImage
+
+
+def handle_image(image: NDArray):
+    result = ocr(image)
+    ret, encoded_image = cv2.imencode(".png", image)
+
+    if not ret:
+        raise Exception("Unknown error.")
+
+    with Image(blob=encoded_image) as img:
+        texts: list[str] = []
+        rects: list[tuple[int, int, int, int]] = []
+        with Drawing() as draw:
+            draw.fill_color = Color("white")
+            for shape, text in result:
+                x, y, xmax, ymax = map(int, shape.bounds)
+                w = xmax - x
+                h = ymax - y
+                draw.rectangle(x, y, xmax, ymax)
+                rects.append((x, y, w, h))
+                texts.append(text)
+            draw.draw(img)
+        font = Font(FONT_PATH)
+        translator = GoogleTranslator(source="auto", target="vi")
+        texts = translator.translate_batch(texts)
+        for (x, y, w, h), text in zip(rects, texts):
+            img.caption(text, x, y, w, h, font=font, gravity="center")
+
+        return PILImage.open(io.BytesIO(img.make_blob("png")))
+
+
+with gr.Blocks() as demo:
+    with gr.Column():
+        image = gr.Image()
+        output = gr.Image()
+        image.upload(handle_image, inputs=image, outputs=output)
+
+if __name__ == "__main__":
+    demo.launch(server_name='0.0.0.0')

tests/test_ocr.py

@@ -0,0 +1,6 @@
+from app import ocr
+from pathlib import Path
+
+def test_detection():
+    result = ocr(str(Path(__file__).parent / 'demo.jpg'))
+    assert len(result) == 11