Spaces:
Running on Zero
Running on Zero
| """CRAFT character/word detection and PP-OCRv6 CTC recognition. | |
| Models are installed explicitly; inference never downloads files or calls an OCR service. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import math | |
| import os | |
| from functools import lru_cache | |
| from pathlib import Path | |
| import numpy as np | |
| from PIL import Image | |
| def model_files(root: Path) -> tuple[Path, Path, Path]: | |
| paths = ( | |
| root / "craft_mlt_25k.pth", | |
| root / "PP-OCRv6_medium_rec_onnx/inference.onnx", | |
| root / "PP-OCRv6_medium_rec_onnx/inference.yml", | |
| ) | |
| missing = [str(p.name) for p in paths if not p.is_file()] | |
| if missing: | |
| raise ValueError( | |
| "Install the optional OCR models in ComfyUI/models/pixtext_ocr. " | |
| "Download links and paths are in the English workflow's OCR setup note. Missing: " | |
| + ", ".join(missing) | |
| ) | |
| return paths | |
| def model_identity(root: Path): | |
| return tuple((str(p), p.stat().st_size, p.stat().st_mtime_ns) for p in model_files(root)) | |
| def decode_ctc(probabilities: np.ndarray, vocabulary: list[str], valid_steps: int): | |
| """Collapse repeats before removing blanks, preserving repeated letters across blanks.""" | |
| sequence = probabilities[:valid_steps] | |
| ids, scores = sequence.argmax(-1), sequence.max(-1) | |
| text, confidence, previous = [], [], -1 | |
| for index, score in zip(ids, scores, strict=True): | |
| index = int(index) | |
| if index and index != previous: | |
| if index >= len(vocabulary): | |
| raise ValueError("OCR dictionary does not match the recognition model") | |
| text.append(vocabulary[index]) | |
| confidence.append(float(score)) | |
| previous = index | |
| return "".join(text), float(np.mean(confidence)) if confidence else 0.0 | |
| def crop_quad(rgb: np.ndarray, quad: np.ndarray): | |
| import cv2 | |
| width = max(2, round(max(np.linalg.norm(quad[1] - quad[0]), np.linalg.norm(quad[2] - quad[3])))) | |
| height = max( | |
| 2, round(max(np.linalg.norm(quad[3] - quad[0]), np.linalg.norm(quad[2] - quad[1]))) | |
| ) | |
| destination = np.float32([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]]) | |
| matrix = cv2.getPerspectiveTransform(quad.astype(np.float32), destination) | |
| return cv2.warpPerspective(rgb, matrix, (width, height)) | |
| def has_han(text: str): | |
| return any(0x3400 <= ord(c) <= 0x9FFF or 0xF900 <= ord(c) <= 0xFAFF | |
| or 0x20000 <= ord(c) <= 0x323AF for c in text) | |
| class WordOCR: | |
| def __init__(self, identity): | |
| try: | |
| import onnxruntime as ort | |
| import torch | |
| import yaml | |
| from easyocr.craft import CRAFT | |
| from easyocr.detection import copyStateDict | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "Install requirements-ocr.txt into ComfyUI's Python environment to use automatic OCR." | |
| ) from exc | |
| craft_path, rec_path, dictionary = [Path(row[0]) for row in identity] | |
| self.detector = CRAFT() | |
| self.detector.load_state_dict( | |
| copyStateDict(torch.load(craft_path, map_location="cpu", weights_only=True)) | |
| ) | |
| self.detector.eval() | |
| options = ort.SessionOptions() | |
| options.intra_op_num_threads = int(os.environ.get("PIXTEXT_CPU_THREADS", "4")) | |
| options.inter_op_num_threads = 1 | |
| self.recognizer = ort.InferenceSession( | |
| str(rec_path), sess_options=options, providers=["CPUExecutionProvider"] | |
| ) | |
| config = yaml.safe_load(dictionary.read_text(encoding="utf-8")) | |
| chars = config["PostProcess"]["character_dict"] | |
| self.vocabulary = ( | |
| ["blank"] + chars + ([" "] if config["PostProcess"].get("use_space_char", True) else []) | |
| ) | |
| self.input_name = self.recognizer.get_inputs()[0].name | |
| def recognize(self, crops): | |
| import cv2 | |
| results = [] | |
| # Small batches bound RAM for the large multilingual output vocabulary. | |
| for start in range(0, len(crops), 8): | |
| group = crops[start : start + 8] | |
| widths = [min(3200, max(1, math.ceil(48 * c.shape[1] / c.shape[0]))) for c in group] | |
| width = max(320, max(widths)) | |
| batch = np.zeros((len(group), 3, 48, width), dtype=np.float32) | |
| for i, (crop, valid_width) in enumerate(zip(group, widths, strict=True)): | |
| # Official PP-OCR preprocessing uses BGR, not RGB. | |
| resized = cv2.resize(crop[:, :, ::-1], (valid_width, 48)) | |
| batch[i, :, :, :valid_width] = ( | |
| (resized.astype(np.float32) / 255 - 0.5) / 0.5 | |
| ).transpose(2, 0, 1) | |
| predictions = self.recognizer.run(None, {self.input_name: batch})[0] | |
| if predictions.shape[-1] != len(self.vocabulary): | |
| raise ValueError("OCR model output and dictionary have different vocabulary sizes") | |
| for pred, valid_width in zip(predictions, widths, strict=True): | |
| results.append( | |
| decode_ctc( | |
| pred, | |
| self.vocabulary, | |
| min(len(pred), math.ceil(len(pred) * valid_width / width)), | |
| ) | |
| ) | |
| return results | |
| def detect_quads(self, rgb: np.ndarray, text_threshold: float, link_threshold: float): | |
| import cv2 | |
| import torch | |
| from easyocr.craft_utils import adjustResultCoordinates, getDetBoxes | |
| from easyocr.imgproc import normalizeMeanVariance, resize_aspect_ratio | |
| resized, ratio, _ = resize_aspect_ratio( | |
| rgb, 1280, interpolation=cv2.INTER_LINEAR, mag_ratio=1.5 | |
| ) | |
| tensor = torch.from_numpy(normalizeMeanVariance(resized)).permute(2, 0, 1).unsqueeze(0) | |
| with torch.inference_mode(): | |
| prediction, _ = self.detector(tensor) | |
| text_map = prediction[0, :, :, 0].cpu().numpy() | |
| link_map = prediction[0, :, :, 1].cpu().numpy() | |
| words, _, _ = getDetBoxes(text_map, link_map, text_threshold, link_threshold, 0.4) | |
| # Region scores locate glyphs. Affinity scores are only used to join them into words. | |
| characters, _, _ = getDetBoxes( | |
| text_map, np.zeros_like(link_map), text_threshold, 1.0, 0.4 | |
| ) | |
| return (list(adjustResultCoordinates(words, 1 / ratio, 1 / ratio)), | |
| list(adjustResultCoordinates(characters, 1 / ratio, 1 / ratio))) | |
| def detect(self, image: Image.Image, text_threshold: float, link_threshold: float): | |
| import cv2 | |
| rgb = np.array(image.convert("RGB")) | |
| words, characters = self.detect_quads(rgb, text_threshold, link_threshold) | |
| if len(words) > 256 or len(characters) > 2048: | |
| raise ValueError("Too much text detected. Crop the image or raise the text threshold.") | |
| word_results = self.recognize([crop_quad(rgb, q) for q in words]) | |
| assigned = {i: [] for i in range(len(words))} | |
| for quad in characters: | |
| center = tuple(float(x) for x in quad.mean(axis=0)) | |
| parents = [i for i, word in enumerate(words) | |
| if cv2.pointPolygonTest(word.astype(np.float32), center, False) >= 0] | |
| if parents: | |
| parent = min(parents, key=lambda i: abs(cv2.contourArea(words[i].astype(np.float32)))) | |
| assigned[parent].append(quad) | |
| proposals, character_tasks = [], [] | |
| for i, (word, (text, score)) in enumerate(zip(words, word_results, strict=True)): | |
| inner = assigned[i] | |
| if has_han(text) or (not text.strip() and inner): | |
| # Sort along the detected text direction, including vertical and rotated text. | |
| across, down = word[1] - word[0], word[3] - word[0] | |
| axis = across if np.linalg.norm(across) >= np.linalg.norm(down) else down | |
| inner.sort(key=lambda q: float(np.dot(q.mean(axis=0), axis))) | |
| for quad in inner: | |
| character_tasks.append((i, quad)) | |
| elif text.strip(): | |
| proposals.append((word, text, score, "word")) | |
| recognized = self.recognize([crop_quad(rgb, quad) for _, quad in character_tasks]) | |
| grouped = {} | |
| for (parent, quad), (text, score) in zip(character_tasks, recognized, strict=True): | |
| # Keep an unreadable detected glyph editable rather than assigning a phrase to it. | |
| text = text.strip() | |
| grouped.setdefault(parent, []).append((quad, text if len(text) == 1 else "", | |
| score if len(text) == 1 else 0.0)) | |
| for parent, glyphs in grouped.items(): | |
| parent_text = word_results[parent][0] | |
| mixed = not parent_text.strip() or any(c.isascii() and c.isalnum() for c in parent_text) | |
| latin = [] | |
| def flush_latin(): | |
| if latin: | |
| points = np.concatenate([r[0] for r in latin]).astype(np.float32) | |
| quad = cv2.boxPoints(cv2.minAreaRect(points)) | |
| proposals.append((quad, "".join(r[1] for r in latin), | |
| float(np.mean([r[2] for r in latin])), "word")) | |
| latin.clear() | |
| for quad, text, score in glyphs: | |
| if mixed and text and text.isascii() and text.isalnum(): | |
| latin.append((quad, text, score)) | |
| else: | |
| flush_latin() | |
| proposals.append((quad, text, score, "character")) | |
| flush_latin() | |
| if len(proposals) > 256: | |
| raise ValueError("More than 256 regions detected. Crop the image or raise the text threshold.") | |
| rows = [] | |
| for quad, text, score, level in proposals: | |
| low = np.maximum(0, quad.min(axis=0)) | |
| high = np.minimum([image.width, image.height], quad.max(axis=0)) | |
| if np.any(high <= low): | |
| continue | |
| bbox = [round(float(x), 2) for x in [*low, *high]] | |
| key = json.dumps([bbox, text], ensure_ascii=False) | |
| rows.append( | |
| { | |
| "id": "ocr-" + hashlib.sha256(key.encode()).hexdigest()[:16], | |
| "bbox": bbox, | |
| "text": text, | |
| "recognized_text": text, | |
| "confidence": round(score, 4), | |
| "enabled": False, | |
| "origin": "ocr", | |
| "level": level, | |
| } | |
| ) | |
| rows.sort(key=lambda row: (round(row["bbox"][1] / 12), row["bbox"][0])) | |
| return rows | |
| def load_ocr(identity): | |
| return WordOCR(identity) | |
| def detect_words(image: Image.Image, root: Path, text_threshold=0.5, link_threshold=0.2): | |
| identity = model_identity(root) | |
| rows = load_ocr(identity).detect(image, text_threshold, link_threshold) | |
| source_id = "pixels:" + hashlib.sha256(image.convert("RGB").tobytes()).hexdigest() | |
| document = { | |
| "version": 1, | |
| "source_id": source_id, | |
| "size": list(image.size), | |
| "regions": rows, | |
| "engine": "CRAFT characters/words + PP-OCRv6 medium (CPU)", | |
| } | |
| document["detection_id"] = hashlib.sha256( | |
| json.dumps(document, sort_keys=True, ensure_ascii=False).encode() | |
| ).hexdigest() | |
| return document | |