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")]