commit 578cc3cbf2a92d4fbcc412bb3dc71f60a6369aa8 Author: malvm Date: Wed Jun 24 13:01:26 2026 +0400 first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..5e68fa8 --- /dev/null +++ b/.gitignore @@ -0,0 +1,13 @@ +# Локальная модель TrOCR (большие веса — не коммитим) +/trocr/ + +# Данные разметки и кэш модели +/data/images/ +/data/labels*.tsv +/data/autolabel_*.json +/data/hf_cache/ + +# Python +__pycache__/ +*.pyc +.venv/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..0ce2c08 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,20 @@ +FROM python:3.11-slim + +WORKDIR /app + +RUN pip install --no-cache-dir flask==3.0.3 gunicorn==22.0.0 +# TrOCR auto-labeling deps (CPU build of torch keeps the image smaller). +RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu +RUN pip install --no-cache-dir transformers pillow sentencepiece + +# Persist the downloaded HuggingFace model in the mounted /data volume. +ENV HF_HOME=/data/hf_cache + +COPY app.py . +COPY templates/ templates/ + +EXPOSE 5000 + +# Single worker + threads: the TrOCR model is loaded once per process and the +# auto-label job tracks progress in-process, so we keep one worker. +CMD ["gunicorn", "--bind", "0.0.0.0:5000", "--workers", "1", "--threads", "4", "--timeout", "120", "app:app"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..5420bda --- /dev/null +++ b/README.md @@ -0,0 +1,156 @@ +# OCR Annotator + +Веб-инструмент для разметки кропов текста. Показывает изображение, вы вводите текст, +результат сохраняется в `labels.tsv` — готовый формат для дообучения TrOCR. + +## Быстрый старт через Docker + +### 1. Установите Docker и Docker Compose (если ещё нет) + +```bash +# Ubuntu / Debian +sudo apt update +sudo apt install -y docker.io docker-compose-plugin + +# Добавьте себя в группу docker (чтобы не писать sudo) +sudo usermod -aG docker $USER +newgrp docker +``` + +### 2. Скачайте / разархивируйте проект + +```bash +# Перейдите в папку проекта +cd ocr-annotator +``` + +### 3. Положите ваши кропы в папку data/images + +```bash +mkdir -p data/images + +# Скопируйте ваши кропы: +cp /path/to/your/crops/*.jpg data/images/ +# или +cp -r /path/to/your/crops/. data/images/ +``` + +Поддерживаемые форматы: `jpg`, `jpeg`, `png`, `bmp`, `tiff`, `webp` + +### 4. Запустите + +```bash +docker compose up --build +``` + +Откройте браузер: **http://localhost:5000** + +При следующих запусках (образ уже собран): +```bash +docker compose up +``` + +Остановить: +```bash +docker compose down +``` + +--- + +## 🤖 Авто-разметка (TrOCR) + +В шапке рабочего экрана есть кнопка **🤖 Авто-разметка**. Она прогоняет все +**неразмеченные** картинки текущей подпапки через модель TrOCR и сохраняет +предсказания с пометкой «авто» — то есть **требуется ручная перепроверка**. + +- Авто-метки в списке файлов и в счётчике отмечены оранжевым и значком 🤖. +- Фильтр **🤖 Проверить** показывает только авто-метки, ожидающие проверки. +- Как только вы вручную сохраняете такую метку (Enter / «Сохранить»), пометка + «авто» снимается — метка становится проверенной. +- Прогресс показывается в статус-баре; разметка идёт в фоне, можно продолжать + проверять уже готовые картинки. + +Модель задаётся переменной `TROCR_MODEL` (по умолчанию `raxtemur/trocr-base-ru` +— TrOCR, дообученная на русский). Можно указать любую HuggingFace-модель типа +`VisionEncoderDecoder`, например `microsoft/trocr-base-printed`. При первом +запуске модель скачивается из HuggingFace и кэшируется в `data/hf_cache` +(переменная `HF_HOME`), поэтому повторные запуски быстрые. + +> В экспортируемый `labels.tsv` пометка «авто» **не попадает** — там остаётся +> чистый формат `имя_файлатекст` для дообучения TrOCR. + +--- + +## Горячие клавиши + +| Клавиша | Действие | +|---------------|-----------------------------| +| `Enter` | Сохранить и перейти дальше | +| `Tab` | Пропустить (без сохранения) | +| `Alt + →` | Следующее изображение | +| `Alt + ←` | Предыдущее изображение | + +--- + +## Результат + +Файл `data/labels.tsv` — табуляция-разделённый файл: + +``` +crop_001.jpg Иванов И.И. +crop_002.jpg ул. Ленина, д. 5 +crop_003.jpg СЧЁТ-ФАКТУРА +``` + +Скачать через интерфейс: кнопка **⬇ Экспорт TSV**, или забрать прямо из папки `data/labels.tsv`. + +--- + +## Использование для дообучения TrOCR + +```python +from datasets import Dataset +from PIL import Image +import pandas as pd + +# Читаем разметку +df = pd.read_csv("data/labels.tsv", sep="\t", header=None, names=["file_name", "text"]) +df["image"] = df["file_name"].apply(lambda f: Image.open(f"data/images/{f}").convert("RGB")) + +dataset = Dataset.from_pandas(df[["image", "text"]]) +# Далее — стандартный fine-tuning TrOCR через HuggingFace Trainer +``` + +--- + +## Запуск без Docker (если нужно) + +```bash +pip install flask gunicorn +# для кнопки «🤖 Авто-разметка» дополнительно: +pip install torch transformers pillow sentencepiece + +export IMAGES_DIR=./data/images +export OUTPUT_FILE=./data/labels.tsv +export TROCR_MODEL=raxtemur/trocr-base-ru # необязательно + +python app.py +# или через gunicorn: +gunicorn --bind 0.0.0.0:5000 app:app +``` + +--- + +## Структура проекта + +``` +ocr-annotator/ +├── app.py # Flask backend +├── Dockerfile +├── docker-compose.yml +├── templates/ +│ └── index.html # UI +└── data/ + ├── images/ # ← кладите сюда кропы + └── labels.tsv # ← сюда пишется разметка +``` diff --git a/app.py b/app.py new file mode 100644 index 0000000..b38177f --- /dev/null +++ b/app.py @@ -0,0 +1,373 @@ +import os +import csv +import io +import json +import zipfile +import threading +from pathlib import Path +from flask import Flask, render_template, request, jsonify, send_file, send_from_directory + +app = Flask(__name__) + +IMAGES_DIR = Path(os.environ.get("IMAGES_DIR", "/data/images")) +DATA_DIR = Path(os.environ.get("OUTPUT_FILE", "/data/labels.tsv")).parent + +SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp"} + +# TrOCR model used for automatic pre-labeling (Russian-finetuned by default). +TROCR_MODEL = os.environ.get("TROCR_MODEL", "raxtemur/trocr-base-ru") + +_locks = {} +_locks_mutex = threading.Lock() + +# Lazily-loaded TrOCR model, shared across requests in the worker process. +_model_cache = {} +_model_mutex = threading.Lock() + + +def get_lock(subfolder): + with _locks_mutex: + if subfolder not in _locks: + _locks[subfolder] = threading.Lock() + return _locks[subfolder] + + +def get_subfolders(): + if not IMAGES_DIR.exists(): + return [] + subs = [] + for d in sorted(IMAGES_DIR.iterdir()): + if d.is_dir(): + has_images = any( + f.suffix.lower() in SUPPORTED_EXTENSIONS for f in d.iterdir() if f.is_file() + ) + if has_images: + subs.append(d.name) + return subs + + +def get_images_in(subfolder): + folder = IMAGES_DIR / subfolder + if not folder.exists(): + return [] + return sorted([ + f.name for f in folder.iterdir() + if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS + ]) + + +def labels_path(subfolder): + return DATA_DIR / f"labels_{subfolder}.tsv" + + +def load_labels(subfolder): + """Return {filename: {"text": str, "auto": bool}}. + + TSV format: filenametext["auto"]. A third column "auto" marks + labels produced by the automatic pre-labeling pass (need a human recheck). + Rows with two columns are treated as human-verified for backward compat. + """ + path = labels_path(subfolder) + labels = {} + if path.exists(): + with open(path, "r", encoding="utf-8") as f: + for row in csv.reader(f, delimiter="\t"): + if len(row) >= 2: + auto = len(row) >= 3 and row[2] == "auto" + labels[row[0]] = {"text": row[1], "auto": auto} + return labels + + +def save_labels(subfolder, labels): + path = labels_path(subfolder) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8", newline="") as f: + writer = csv.writer(f, delimiter="\t") + for filename, entry in sorted(labels.items()): + row = [filename, entry["text"]] + if entry.get("auto"): + row.append("auto") + writer.writerow(row) + + +@app.route("/") +def index(): + subfolders = get_subfolders() + return render_template("index.html", subfolders=subfolders) + + +@app.route("/api/subfolders") +def api_subfolders(): + subs = get_subfolders() + result = [] + for s in subs: + images = get_images_in(s) + labels = load_labels(s) + labeled = sum(1 for img in images if img in labels) + auto = sum(1 for img in images if labels.get(img, {}).get("auto")) + result.append({"name": s, "total": len(images), "labeled": labeled, "auto": auto}) + return jsonify(result) + + +@app.route("/api/images/") +def api_images(subfolder): + images = get_images_in(subfolder) + labels = load_labels(subfolder) + return jsonify([ + { + "filename": img, + "labeled": img in labels, + "text": labels.get(img, {}).get("text", ""), + "auto": labels.get(img, {}).get("auto", False), + } + for img in images + ]) + + +@app.route("/api/label/", methods=["POST"]) +def api_label(subfolder): + data = request.json + filename = data.get("filename") + text = (data.get("text") or "").strip() + + if not filename: + return jsonify({"error": "filename required"}), 400 + + images = get_images_in(subfolder) + if filename not in images: + return jsonify({"error": "image not found"}), 404 + + lock = get_lock(subfolder) + with lock: + labels = load_labels(subfolder) + if text == "" and filename in labels: + del labels[filename] + else: + # A manual save always counts as human-verified (clears the auto flag). + labels[filename] = {"text": text, "auto": False} + save_labels(subfolder, labels) + + images_after = get_images_in(subfolder) + labels_after = load_labels(subfolder) + labeled = sum(1 for img in images_after if img in labels_after) + return jsonify({"ok": True, "labeled": labeled, "total": len(images_after)}) + + +@app.route("/api/delete/", methods=["POST"]) +def api_delete(subfolder): + data = request.json + filename = data.get("filename") + + if not filename: + return jsonify({"error": "filename required"}), 400 + + # Validate: must be inside the subfolder, no path traversal + image_path = (IMAGES_DIR / subfolder / filename).resolve() + allowed_root = (IMAGES_DIR / subfolder).resolve() + if not str(image_path).startswith(str(allowed_root)): + return jsonify({"error": "invalid path"}), 400 + + if not image_path.exists(): + return jsonify({"error": "file not found"}), 404 + + lock = get_lock(subfolder) + with lock: + # Remove from disk + image_path.unlink() + # Remove label if exists + labels = load_labels(subfolder) + if filename in labels: + del labels[filename] + save_labels(subfolder, labels) + + images_after = get_images_in(subfolder) + labels_after = load_labels(subfolder) + labeled = sum(1 for img in images_after if img in labels_after) + return jsonify({"ok": True, "labeled": labeled, "total": len(images_after)}) + + +@app.route("/images//") +def serve_image(subfolder, filename): + return send_from_directory(IMAGES_DIR / subfolder, filename) + + +@app.route("/api/export") +def api_export(): + """Export a ZIP: labeled images + labels.tsv""" + subfolders = get_subfolders() + rows = [] + image_paths = [] # (arc_name, fs_path) + + for sub in subfolders: + labels = load_labels(sub) + for filename, entry in sorted(labels.items()): + text = entry["text"] + fs_path = IMAGES_DIR / sub / filename + if fs_path.exists(): + arc_name = f"{sub}/{filename}" + rows.append((arc_name, text)) + image_paths.append((arc_name, fs_path)) + + if not rows: + return jsonify({"error": "No labeled images yet"}), 404 + + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: + # Write TSV + tsv_buf = io.StringIO() + writer = csv.writer(tsv_buf, delimiter="\t") + for row in rows: + writer.writerow(row) + zf.writestr("labels.tsv", tsv_buf.getvalue().encode("utf-8")) + + # Write images + for arc_name, fs_path in image_paths: + zf.write(fs_path, arc_name) + + buf.seek(0) + return send_file( + buf, + mimetype="application/zip", + as_attachment=True, + download_name="dataset.zip", + ) + + +@app.route("/api/stats") +def api_stats(): + subs = get_subfolders() + total = labeled = auto = 0 + for s in subs: + imgs = get_images_in(s) + lbls = load_labels(s) + total += len(imgs) + labeled += sum(1 for img in imgs if img in lbls) + auto += sum(1 for img in imgs if lbls.get(img, {}).get("auto")) + return jsonify({"total": total, "labeled": labeled, "remaining": total - labeled, "auto": auto}) + + +# ── Automatic pre-labeling with TrOCR ─────────────────────────────────────── + +def autolabel_status_path(subfolder): + return DATA_DIR / f"autolabel_{subfolder}.json" + + +def read_autolabel_status(subfolder): + """Progress is persisted to a file so any gunicorn worker can read it.""" + path = autolabel_status_path(subfolder) + if path.exists(): + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return None + return None + + +def write_autolabel_status(subfolder, status): + path = autolabel_status_path(subfolder) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(status), encoding="utf-8") + + +def get_trocr(): + """Load (once) and return the TrOCR processor + model. Heavy imports are + done lazily so the app still runs in environments without torch installed.""" + with _model_mutex: + if "model" not in _model_cache: + import torch + from transformers import TrOCRProcessor, VisionEncoderDecoderModel + + processor = TrOCRProcessor.from_pretrained(TROCR_MODEL) + model = VisionEncoderDecoderModel.from_pretrained(TROCR_MODEL) + device = "cuda" if torch.cuda.is_available() else "cpu" + model.to(device) + model.eval() + _model_cache.update( + {"processor": processor, "model": model, "device": device, "torch": torch} + ) + return ( + _model_cache["processor"], + _model_cache["model"], + _model_cache["device"], + _model_cache["torch"], + ) + + +def run_autolabel(subfolder): + """Background job: run TrOCR over every unlabeled image in the subfolder and + store the predictions flagged as `auto` (i.e. needing a human recheck).""" + try: + images = get_images_in(subfolder) + existing = load_labels(subfolder) + todo = [img for img in images if img not in existing] + write_autolabel_status( + subfolder, {"state": "running", "total": len(todo), "done": 0, "error": None} + ) + if not todo: + write_autolabel_status( + subfolder, {"state": "done", "total": 0, "done": 0, "error": None} + ) + return + + from PIL import Image + + processor, model, device, torch = get_trocr() + lock = get_lock(subfolder) + done = 0 + for img in todo: + try: + fs_path = IMAGES_DIR / subfolder / img + image = Image.open(fs_path).convert("RGB") + pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device) + with torch.no_grad(): + generated_ids = model.generate(pixel_values, max_length=64) + text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip() + except Exception: + text = "" + + # Save incrementally and never clobber a label a human added meanwhile. + with lock: + current = load_labels(subfolder) + if img not in current: + current[img] = {"text": text, "auto": True} + save_labels(subfolder, current) + + done += 1 + write_autolabel_status( + subfolder, + {"state": "running", "total": len(todo), "done": done, "error": None}, + ) + + write_autolabel_status( + subfolder, {"state": "done", "total": len(todo), "done": done, "error": None} + ) + except Exception as e: + write_autolabel_status( + subfolder, {"state": "error", "total": 0, "done": 0, "error": str(e)} + ) + + +@app.route("/api/autolabel/", methods=["POST"]) +def api_autolabel(subfolder): + if subfolder not in get_subfolders(): + return jsonify({"error": "unknown subfolder"}), 404 + + status = read_autolabel_status(subfolder) + if status and status.get("state") == "running": + return jsonify({"error": "already running", "status": status}), 409 + + write_autolabel_status( + subfolder, {"state": "running", "total": 0, "done": 0, "error": None} + ) + threading.Thread(target=run_autolabel, args=(subfolder,), daemon=True).start() + return jsonify({"ok": True}) + + +@app.route("/api/autolabel//status") +def api_autolabel_status(subfolder): + status = read_autolabel_status(subfolder) or {"state": "idle"} + return jsonify(status) + + +if __name__ == "__main__": + app.run(host="0.0.0.0", port=5000, debug=False) diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..85eb0bc --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,16 @@ +services: + annotator: + build: . + ports: + - "5010:5000" + volumes: + - ./data/images:/data/images # read-write: нужно для удаления файлов + - ./data:/data # сюда пишется labels_*.tsv + - ./trocr:/models/trocr:ro # локальная (оффлайн) модель TrOCR + environment: + - IMAGES_DIR=/data/images + - OUTPUT_FILE=/data/labels.tsv + # Локальная модель для кнопки «🤖 Авто-разметка» (путь внутри контейнера). + - TROCR_MODEL=/models/trocr + - HF_HOME=/data/hf_cache + restart: unless-stopped diff --git a/templates/index.html b/templates/index.html new file mode 100644 index 0000000..6b3d578 --- /dev/null +++ b/templates/index.html @@ -0,0 +1,741 @@ + + + + + +OCR Annotator + + + + + + + +
+
+ +
Выберите папку для разметки
+
+
+
Загрузка...
+
+
+ Каждый пользователь выбирает свою подпапку — разметка сохраняется независимо.
+ При экспорте все подпапки объединяются в ZIP (картинки + labels.tsv). +
+
+ + +
+
+ + +
—
+
+
+
0 / 0 размечено
+
+
+ + + + +
+
+ +
+
+
+ Файлы +
+ + + + +
+
+
+
+ +
+
+ +
+ +
100%
+ + +
+ +
+ +
+
+
🗂
+

Нет изображений

+
+
+ +
+
+ Текст → + + +
+
+
+ + +
+
+
Enter сохранить и далее
+
Tab пропустить
+
Alt+←/→ навигация
+
колесо зум
+
+
+
+
+
+ +
+
Готов
+
—
+
+
+ + + + + +
+ + + +