From 65e4a3554640841954c18934d444abe7f23c285b Mon Sep 17 00:00:00 2001 From: chopratejas Date: Thu, 9 Apr 2026 16:22:18 -0700 Subject: [PATCH] Add OCR routing via RapidOCR + measure actual tokens after compression MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Transcode technique now runs RapidOCR to extract text from images. Falls back to full_low if OCR confidence < 70% or no text detected. Token counting is now done AFTER compression by measuring the actual output — no hardcoded estimates. OCR text counted by char length, resized images re-estimated from new dimensions. RapidOCR added to headroom-ai[image] extra (~15MB ONNX models, ~180ms CPU). 4 new OCR tests (extraction, blank image, confidence threshold, full pipeline). 29 total image compression tests passing. --- headroom/image/compressor.py | 287 ++++++++++++++++++++++---------- pyproject.toml | 3 +- tests/test_image_compression.py | 119 +++++++++++++ 3 files changed, 318 insertions(+), 91 deletions(-) diff --git a/headroom/image/compressor.py b/headroom/image/compressor.py index 71db4e00e..27714afa9 100644 --- a/headroom/image/compressor.py +++ b/headroom/image/compressor.py @@ -241,6 +241,104 @@ class ImageCompressor: tiles_y = (height + 511) // 512 return int(85 * tiles_x * tiles_y + 170) + def _count_result_tokens( + self, + messages: list[dict[str, Any]], + original_image_data: bytes, + provider: str, + ) -> int: + """Count actual tokens in compressed messages by inspecting the result. + + If the image was replaced with OCR text → count text tokens (~4 chars/token). + If the image was resized → re-estimate from new dimensions. + If detail=low was set → use provider's low-detail cost. + """ + total = 0 + for message in messages: + content = message.get("content") + if not isinstance(content, list): + continue + + for item in content: + if not isinstance(item, dict): + continue + + # OCR replacement: text block with "[OCR from image]" + if item.get("type") == "text" and "[OCR from image]" in item.get("text", ""): + text = item["text"] + total += max(1, len(text) // 4) # ~4 chars per token + continue + + # OpenAI: check if detail was set to "low" + if item.get("type") == "image_url": + detail = item.get("image_url", {}).get("detail", "high") + if detail == "low": + total += 85 # OpenAI's documented low-detail cost + else: + # Re-estimate from the (possibly resized) image + url = item.get("image_url", {}).get("url", "") + if url.startswith("data:"): + match = re.match(r"data:image/[^;]+;base64,(.+)", url) + if match: + data = base64.b64decode(match.group(1)) + total += self._estimate_tokens(data, "high") + + # Anthropic: re-estimate from the (possibly resized) image + elif item.get("type") == "image": + source = item.get("source", {}) + if source.get("type") == "base64": + data = base64.b64decode(source.get("data", "")) + total += self._estimate_tokens(data, "high") + + # Google: re-estimate + elif "inlineData" in item: + data = base64.b64decode(item.get("inlineData", {}).get("data", "")) + total += self._estimate_tokens(data, "high") + + return total if total > 0 else self._estimate_tokens(original_image_data, "high") + + def _ocr_extract(self, image_data: bytes, min_confidence: float = 0.7) -> str | None: + """Extract text from image using RapidOCR. + + Returns extracted text if OCR is confident, None otherwise (fallback to image). + """ + try: + from rapidocr_onnxruntime import RapidOCR + + if not hasattr(self, "_ocr_engine"): + self._ocr_engine = RapidOCR() + + result, _ = self._ocr_engine(image_data) + if not result: + return None + + # Check average confidence + confidences = [line[2] for line in result] + avg_confidence = sum(confidences) / len(confidences) + + if avg_confidence < min_confidence: + logger.debug( + f"OCR confidence too low ({avg_confidence:.0%} < {min_confidence:.0%}), " + f"falling back to image" + ) + return None + + # Combine lines into text + text = "\n".join(line[1] for line in result) + + logger.info( + f"OCR extracted {len(result)} lines ({avg_confidence:.0%} avg confidence, " + f"{len(text)} chars)" + ) + return text + + except ImportError: + logger.debug("rapidocr-onnxruntime not installed, skipping OCR") + return None + except Exception as e: + logger.warning(f"OCR failed: {e}") + return None + def _apply_compression( self, messages: list[dict[str, Any]], @@ -265,87 +363,100 @@ class ImageCompressor: new_content.append(item) continue - # OpenAI format - compare by value since technique may be from trained_router - if item.get("type") == "image_url" and provider == "openai": - if technique.value == "full_low": - # Apply detail="low" - new_item = { - "type": "image_url", - "image_url": { - **item.get("image_url", {}), - "detail": "low", - }, - } - new_content.append(new_item) - elif technique.value == "crop": - # For now, use low detail (TODO: implement actual cropping) - new_item = { - "type": "image_url", - "image_url": { - **item.get("image_url", {}), - "detail": "low", - }, - } - new_content.append(new_item) - elif technique.value == "transcode": - # TODO: Convert to text description - # For now, keep original - new_content.append(item) - else: - new_content.append(item) + # Extract image bytes for OCR (transcode) across all formats + image_bytes_for_ocr: bytes | None = None + is_image_block = False - # Anthropic format - resize image for compression - elif item.get("type") == "image" and provider == "anthropic": - if technique.value in ("full_low", "crop"): - # Resize image to reduce tokens - try: - source = item.get("source", {}) - if source.get("type") == "base64": - original_data = base64.b64decode(source.get("data", "")) - resized_data, media_type = self._resize_image( - original_data, max_dimension=512 - ) - new_item = { - "type": "image", - "source": { - "type": "base64", - "media_type": media_type, - "data": base64.b64encode(resized_data).decode(), - }, - } - new_content.append(new_item) - else: - new_content.append(item) - except Exception as e: - logger.warning(f"Failed to resize Anthropic image: {e}") - new_content.append(item) - else: - new_content.append(item) + if item.get("type") == "image_url": + is_image_block = True + url = item.get("image_url", {}).get("url", "") + if url.startswith("data:"): + match = re.match(r"data:image/[^;]+;base64,(.+)", url) + if match: + image_bytes_for_ocr = base64.b64decode(match.group(1)) + elif item.get("type") == "image": + is_image_block = True + source = item.get("source", {}) + if source.get("type") == "base64": + image_bytes_for_ocr = base64.b64decode(source.get("data", "")) + elif "inlineData" in item: + is_image_block = True + image_bytes_for_ocr = base64.b64decode( + item.get("inlineData", {}).get("data", "") + ) - # Google format - resize image for compression - elif "inlineData" in item and provider == "google": - if technique.value in ("full_low", "crop"): - try: - inline = item.get("inlineData", {}) - original_data = base64.b64decode(inline.get("data", "")) - resized_data, media_type = self._resize_image( - original_data, - max_dimension=768, # Google uses 768x768 tiles - ) - new_item = { - "inlineData": { - "mimeType": media_type, - "data": base64.b64encode(resized_data).decode(), - } + if not is_image_block: + new_content.append(item) + continue + + # --- TRANSCODE: OCR the image and replace with text --- + if technique.value == "transcode" and image_bytes_for_ocr: + extracted = self._ocr_extract(image_bytes_for_ocr) + if extracted: + # Replace image with extracted text + new_content.append( + {"type": "text", "text": f"[OCR from image]\n{extracted}"} + ) + continue + # OCR failed or low confidence — fall through to full_low + logger.debug("OCR fallback: using full_low instead of transcode") + + # --- FULL_LOW / CROP: reduce quality --- + if technique.value in ("full_low", "crop", "transcode"): + if item.get("type") == "image_url" and provider == "openai": + new_content.append( + { + "type": "image_url", + "image_url": { + **item.get("image_url", {}), + "detail": "low", + }, } - new_content.append(new_item) - except Exception as e: - logger.warning(f"Failed to resize Google image: {e}") + ) + elif item.get("type") == "image" and provider == "anthropic": + if image_bytes_for_ocr: + try: + resized_data, media_type = self._resize_image( + image_bytes_for_ocr, max_dimension=512 + ) + new_content.append( + { + "type": "image", + "source": { + "type": "base64", + "media_type": media_type, + "data": base64.b64encode(resized_data).decode(), + }, + } + ) + except Exception as e: + logger.warning(f"Failed to resize image: {e}") + new_content.append(item) + else: + new_content.append(item) + elif "inlineData" in item and provider == "google": + if image_bytes_for_ocr: + try: + resized_data, media_type = self._resize_image( + image_bytes_for_ocr, max_dimension=768 + ) + new_content.append( + { + "inlineData": { + "mimeType": media_type, + "data": base64.b64encode(resized_data).decode(), + } + } + ) + except Exception as e: + logger.warning(f"Failed to resize image: {e}") + new_content.append(item) + else: new_content.append(item) else: new_content.append(item) - else: + # PRESERVE or unknown — keep original new_content.append(item) compressed.append({**message, "content": new_content}) @@ -424,24 +535,21 @@ class ImageCompressor: technique = Technique.PRESERVE confidence = 0.0 - # Calculate tokens - original_tokens = self._estimate_tokens(image_data, "high") + # Count original tokens BEFORE compression + original_tokens = self._estimate_tokens(image_data, "high") + tile_saved - if technique.value == "full_low": - compressed_tokens = 85 # OpenAI low detail - elif technique.value == "preserve": - compressed_tokens = original_tokens - elif technique.value == "crop": - compressed_tokens = 85 # Approximation - elif technique.value == "transcode": - compressed_tokens = 50 # Text description estimate - else: - compressed_tokens = original_tokens + # Step 3: Apply compression technique + compressed_messages = self._apply_compression(messages, technique, provider) - # Store result (include tile savings) + # Count actual tokens AFTER compression by measuring the result. + # If the image was replaced with text (OCR), count text tokens. + # If resized, re-estimate from new dimensions. + compressed_tokens = self._count_result_tokens(compressed_messages, image_data, provider) + + # Store result self.last_result = CompressionResult( technique=technique, - original_tokens=original_tokens + tile_saved, + original_tokens=original_tokens, compressed_tokens=compressed_tokens, confidence=confidence, ) @@ -452,8 +560,7 @@ class ImageCompressor: f"{self.last_result.savings_percent:.0f}% saved)" ) - # Step 3: Apply compression technique - return self._apply_compression(messages, technique, provider) + return compressed_messages # Singleton for convenience diff --git a/pyproject.toml b/pyproject.toml index 7111a6cf5..d8b0e037e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,10 +92,11 @@ relevance = [ "sentence-transformers>=2.2.0", "numpy>=1.24.0", ] -# Image compression (ML-based routing) +# Image compression (ML-based routing + OCR) image = [ "pillow>=10.0.0", "sentencepiece>=0.1.99", # Required by SigLIP tokenizer (SiglipTokenizer) + "rapidocr-onnxruntime>=1.4.0", # ONNX-native OCR for text extraction from images (~15MB models) ] # Report generation reports = [ diff --git a/tests/test_image_compression.py b/tests/test_image_compression.py index 240b8c0d9..7480d1471 100644 --- a/tests/test_image_compression.py +++ b/tests/test_image_compression.py @@ -305,3 +305,122 @@ class TestFullPipeline: compressor = ImageCompressor() msgs = [{"role": "user", "content": "Just text"}] assert not compressor.has_images(msgs) + + +# --------------------------------------------------------------------------- +# OCR routing tests +# --------------------------------------------------------------------------- + + +class TestOcrRouting: + @pytest.fixture(autouse=True) + def _check_ocr(self): + try: + from rapidocr_onnxruntime import RapidOCR # noqa: F401 + except ImportError: + pytest.skip("rapidocr-onnxruntime not installed") + + def _make_text_image(self, lines: list[str], width: int = 800, height: int = 400) -> bytes: + """Create a PNG image with text content.""" + from PIL import Image, ImageDraw + + img = Image.new("RGB", (width, height), "white") + draw = ImageDraw.Draw(img) + y = 30 + for line in lines: + draw.text((30, y), line, fill="black") + y += 40 + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + def test_ocr_extracts_text(self): + """OCR should extract text from a text-heavy image.""" + from headroom.image import ImageCompressor + + compressor = ImageCompressor(use_siglip=False) + image_data = self._make_text_image( + [ + "Error: connection refused", + "at localhost:5432", + ] + ) + text = compressor._ocr_extract(image_data) + assert text is not None + assert len(text) > 10 + # Should contain key words (OCR may have minor errors) + assert "connection" in text.lower() or "error" in text.lower() + + def test_ocr_returns_none_for_blank_image(self): + """OCR should return None for a blank image (no text).""" + from headroom.image import ImageCompressor + + compressor = ImageCompressor(use_siglip=False) + from PIL import Image + + img = Image.new("RGB", (200, 200), "blue") + buf = io.BytesIO() + img.save(buf, format="PNG") + text = compressor._ocr_extract(buf.getvalue()) + assert text is None # No text detected + + def test_ocr_confidence_threshold(self): + """Low-confidence OCR should return None (fallback to image).""" + from headroom.image import ImageCompressor + + compressor = ImageCompressor(use_siglip=False) + # Very noisy image — OCR should have low confidence + import numpy as np + from PIL import Image + + noise = np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8) + img = Image.fromarray(noise) + buf = io.BytesIO() + img.save(buf, format="PNG") + text = compressor._ocr_extract(buf.getvalue(), min_confidence=0.95) + # Noisy image: either None (no text) or low confidence → None + # Either outcome is correct — we don't want to OCR noise + assert text is None or len(text) < 10 + + def test_transcode_replaces_image_with_text(self): + """Full pipeline: transcode technique should replace image with OCR text.""" + from headroom.image import ImageCompressor + from headroom.image.trained_router import Technique + + compressor = ImageCompressor(use_siglip=False) + + # Create message with text-heavy image + image_data = self._make_text_image( + [ + "Traceback (most recent call last):", + " File server.py line 42", + "psycopg2.OperationalError", + ] + ) + b64 = base64.b64encode(image_data).decode() + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What does the error say?"}, + { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{b64}"}, + }, + ], + } + ] + + # Apply transcode directly + result = compressor._apply_compression(messages, Technique.TRANSCODE, "openai") + + # The image block should be replaced with a text block + content = result[0]["content"] + text_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "text"] + + # Should have at least 2 text blocks (original query + OCR output) + assert len(text_blocks) >= 2 + # One should contain OCR output + ocr_blocks = [b for b in text_blocks if "[OCR from image]" in b.get("text", "")] + assert len(ocr_blocks) >= 1