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