test_ocr.py

import fitz  # PyMuPDF
import os

def ocr_file(path: str, language="eng") -> str:
    ext = os.path.splitext(path)[1].lower()

    if ext == ".pdf":
        # Open PDF, take first page
        doc = fitz.open(path)
        page = doc[0]
        tp = page.get_textpage_ocr(language=language)
        return page.get_text("text", textpage=tp)

    elif ext in [".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"]:
        pix = fitz.Pixmap(path)

        # Ensure RGB, no alpha
        if pix.colorspace is None or pix.colorspace.n != 3:
            pix = fitz.Pixmap(fitz.csRGB, pix)
        if pix.alpha:
            pix = fitz.Pixmap(pix, 0)

        # OCR as a one-page PDF in memory
        img_pdf_bytes = pix.pdfocr_tobytes(language=language)
        doc = fitz.open("pdf", img_pdf_bytes)
        return doc[0].get_text("text")

    else:
        raise ValueError(f"Unsupported file type: {ext}")


# Example usage
if __name__ == "__main__":
    text = ocr_file("/Users/nyeung/Downloads/stock_q.png", language="eng")
    print(text)