Experiments / processor_utils.py
credent007's picture
update load_input function which now can take both image and pdf
b17ed36 verified
raw
history blame contribute delete
797 Bytes
from PIL import Image
import fitz
# def load_input(file_path): # past code
# if file_path.lower().endswith(".pdf"):
# doc = fitz.open(file_path)
# page = doc[0]
# pix = page.get_pixmap()
# img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
# return img
# else:
# return Image.open(file_path).convert("RGB")
def load_input(file_path):
if file_path.lower().endswith(".pdf"):
doc = fitz.open(file_path)
images = []
for page in doc:
pix = page.get_pixmap()
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
images.append(img)
return images # 🔥 return list
else:
return [Image.open(file_path).convert("RGB")]