S SmartDocs
シリーズ: Ricky python 85 行 · 更新日 2026-04-07

image_processor.py

Ricky/RAG_project/processors/image_processor.py

"""
Image ingestion with OCR.

When OCR produces no text, the filename is used as a minimal searchable
index entry so image-intent queries can still match by product name.
OCR quality score is stored in metadata for potential downstream filtering.
"""
import os
import logging

from langchain_core.documents import Document

from config.settings import RAW_IMAGE_PATH, OCR_LANGUAGE, TESSERACT_CMD
from processors.source_metadata import build_source_metadata
from processors.pdf_processor import _compute_text_quality  # reuse quality scorer

logger = logging.getLogger(__name__)

try:
    from PIL import Image
except Exception:
    Image = None  # type: ignore

try:
    import pytesseract
except Exception:
    pytesseract = None  # type: ignore

_SUPPORTED_EXT = (".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff")


def process_images() -> list[Document]:
    """Load image files, OCR them, and return as Document objects."""
    documents: list[Document] = []
    if not os.path.exists(RAW_IMAGE_PATH):
        print("Image folder not found")
        return documents

    if Image is None or pytesseract is None:
        print("[WARN] Missing image OCR packages. Install: pip install pillow pytesseract")
        return documents

    if TESSERACT_CMD:
        pytesseract.pytesseract.tesseract_cmd = TESSERACT_CMD

    for filename in os.listdir(RAW_IMAGE_PATH):
        if not filename.lower().endswith(_SUPPORTED_EXT):
            continue
        filepath = os.path.join(RAW_IMAGE_PATH, filename)
        print(f"Processing: {filename}")
        try:
            image = Image.open(filepath)
            text = pytesseract.image_to_string(image, lang=OCR_LANGUAGE)
            cleaned = " ".join(text.split())
            ocr_quality: float | None = None

            if cleaned:
                ocr_quality = _compute_text_quality(cleaned)
            else:
                # Fall back to filename so image-intent queries can still match
                name_only = os.path.splitext(filename)[0]
                cleaned = " ".join(name_only.replace("_", " ").replace("-", " ").split())
                print(f"[WARN] No OCR text extracted from image; using filename as index: {filename}")

            extra = {"ocr_lang": OCR_LANGUAGE, "image_path": filepath}
            if ocr_quality is not None:
                extra["ocr_quality"] = round(ocr_quality, 3)

            documents.append(
                Document(
                    page_content=cleaned,
                    metadata=build_source_metadata(filepath, "image", **extra),
                )
            )
        except Exception as e:
            logger.warning("Skipping unprocessable image: %s (%s)", filename, e)
            print(f"[WARN] Skipping unprocessable image: {filename} ({e})")
            continue

    print(f"Image OCR complete: {len(documents)} images")
    return documents


if __name__ == "__main__":
    docs = process_images()

関連記事