src/services/pdf_processor.py

import fitz


def process_pdf(
    pdf_bytes,
    mode
):
    """
    Returns list of images
    extracted from PDF
    """

    doc = fitz.open(
        stream=pdf_bytes,
        filetype="pdf"
    )

    extracted = []

    if mode == "Extract Embedded Images":

        counter = 1

        for page_num in range(
            len(doc)
        ):

            page = doc[page_num]

            images = (
                page.get_images(
                    full=True
                )
            )

            for img in images:

                xref = img[0]

                base = (
                    doc.extract_image(
                        xref
                    )
                )

                img_bytes = (
                    base["image"]
                )

                ext = (
                    base["ext"]
                )

                if len(
                    img_bytes
                ) < 1024:

                    continue

                extracted.append({

                    "name":
                    f"Extracted Design {counter}",

                    "bytes":
                    img_bytes,

                    "mime_type":
                    f"image/{ext}",

                    "analysis":
                    None,

                    "params":
                    None
                })

                counter += 1

    if (
        mode == "Render Entire Pages"
        or
        not extracted
    ):

        extracted = []

        for page_num in range(
            len(doc)
        ):

            page = doc[page_num]

            pix = (
                page.get_pixmap(
                    dpi=200
                )
            )

            extracted.append({

                "name":
                f"Page {page_num+1}",

                "bytes":
                pix.tobytes("png"),

                "mime_type":
                "image/png",

                "analysis":
                None,

                "params":
                None
            })

    return extracted