mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Add OCR routing via RapidOCR + measure actual tokens after compression
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.
This commit is contained in:
chopratejas
parent
7984bd2389
commit
65e4a35546
3 changed files with 318 additions and 91 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue