"""
Qwen-Image-Layered Backend Implementation
Integrates Qwen-Image-Layered for advanced image decomposition and layered OCR
"""
import logging
from pathlib import Path
from typing import Any
import torch
from PIL import Image
from ..core.backend_manager import OCRBackend
from ..core.config import OCRConfig
try:
from diffusers import QwenImageLayeredPipeline
DIFFUSERS_AVAILABLE = True
except ImportError:
DIFFUSERS_AVAILABLE = False
logger = logging.getLogger(__name__)
class QwenLayeredBackend(OCRBackend):
"""Qwen-Image-Layered backend for image decomposition and layered OCR"""
def __init__(self, config: OCRConfig):
super().__init__("qwen-layered", config)
self.pipeline = None
self.device = getattr(config, "ocr_device", None) or (
"cuda" if torch.cuda.is_available() else "cpu"
)
self.model_name = "Qwen/Qwen-Image-Layered"
self.cache_dir = Path(
getattr(config, "ocr_cache_dir", None) or str(Path.home() / ".cache" / "qwen-layered")
)
self.cache_dir.mkdir(parents=True, exist_ok=True)
def is_available(self) -> bool:
"""Check if Qwen-Image-Layered is available"""
if not DIFFUSERS_AVAILABLE:
return False
try:
from huggingface_hub import model_info
model_info(self.model_name)
return True
except Exception as e:
logger.warning(f"Qwen-Image-Layered model check failed: {e}")
return False
async def load_model(self) -> bool:
"""Load the Qwen-Image-Layered model"""
if not DIFFUSERS_AVAILABLE:
logger.error("Diffusers not available for Qwen-Image-Layered")
return False
try:
logger.info(f"Loading Qwen-Image-Layered model on {self.device}")
# Load the layered image pipeline
self.pipeline = QwenImageLayeredPipeline.from_pretrained(
self.model_name,
cache_dir=str(self.cache_dir),
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
)
if self.device == "cuda":
self.pipeline = self.pipeline.to(self.device)
logger.info("Qwen-Image-Layered model loaded successfully")
return True
except Exception as e:
logger.error(f"Failed to load Qwen-Image-Layered model: {e}")
return False
async def process_document(
self,
image_path: str,
ocr_mode: str = "text",
region: list[int] | None = None,
) -> dict[str, Any]:
"""Process document with Qwen-Image-Layered decomposition"""
if not self.pipeline:
raise RuntimeError("Qwen-Image-Layered model not loaded")
try:
# Load and preprocess image
image = Image.open(image_path).convert("RGB")
# Apply region cropping if specified
if region and len(region) == 4:
x1, y1, x2, y2 = region
image = image.crop((x1, y1, x2, y2))
# Decompose image into layers
layers = await self._decompose_image(image)
# Process layers for OCR
ocr_results = await self._process_layers_for_ocr(layers, ocr_mode)
return ocr_results
except Exception as e:
logger.error(f"Qwen-Image-Layered processing failed: {e}")
raise RuntimeError(f"OCR processing failed: {str(e)}")
async def _decompose_image(self, image: Image.Image) -> list[Image.Image]:
"""Decompose image into layers using Qwen-Image-Layered"""
try:
# Generate layered decomposition
# This is a simplified implementation - actual Qwen-Image-Layered
# would decompose into semantic layers
# Placeholder for actual Qwen-Image-Layered processing
# result = self.pipeline(
# prompt=prompt,
# image=image,
# num_inference_steps=20,
# guidance_scale=7.5
# )
# For now, return original image as single layer
# Actual implementation would return multiple RGBA layers
return [image]
except Exception as e:
logger.error(f"Image decomposition failed: {e}")
return [image] # Fallback to original image
async def _process_layers_for_ocr(
self, layers: list[Image.Image], ocr_mode: str
) -> dict[str, Any]:
"""Process decomposed layers for OCR"""
# This would use another OCR backend to process the layers
# For now, return placeholder results
combined_text = ""
regions = []
confidence = 0.85
for i, layer in enumerate(layers):
# Process each layer
# In practice, this would use a text-focused OCR backend
layer_text = f"Layer {i + 1} content" # Placeholder
combined_text += layer_text + "\n"
regions.append(
{
"bbox": [0, 0, layer.width, layer.height],
"text": layer_text,
"confidence": confidence,
"layer": i,
}
)
if ocr_mode == "text":
return {
"text": combined_text.strip(),
"backend": "qwen-layered",
"confidence": confidence,
"regions": [],
}
elif ocr_mode == "format":
return {
"text": combined_text.strip(),
"backend": "qwen-layered",
"confidence": confidence,
"structured": self._create_layered_structure(layers),
"regions": [],
}
else: # fine-grained
return {
"text": combined_text.strip(),
"backend": "qwen-layered",
"confidence": confidence,
"regions": regions,
"structured": self._create_layered_structure(layers),
}
def _create_layered_structure(self, layers: list[Image.Image]) -> dict[str, Any]:
"""Create structured output from layers"""
return {
"layers": len(layers),
"layer_info": [
{
"index": i,
"dimensions": f"{layer.width}x{layer.height}",
"type": "unknown", # Would be determined by decomposition
}
for i, layer in enumerate(layers)
],
"decomposition_method": "qwen-layered",
}
def get_capabilities(self) -> dict[str, Any]:
"""Get backend capabilities"""
return {
"name": "Qwen-Image-Layered",
"available": self.is_available(),
"modes": ["text", "format", "fine-grained"],
"languages": ["en", "zh", "multilingual"],
"gpu_support": True,
"strengths": ["layer_decomposition", "mixed_content", "complex_documents"],
"limitations": ["experimental", "resource_intensive"],
"model_size": "~7GB",
}
