image from pdf

app.py

@@ -10,6 +10,7 @@ from PyPDF2 import PdfReader
 from threading import Thread
 from flask import Flask, request, jsonify
 import io
+import fitz
 
 # Check if we're running in a Hugging Face Space and if SPACES_ZERO_GPU is enabled
 # IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
@@ -64,17 +65,56 @@ processor = AutoProcessor.from_pretrained(model_name, use_auth_token=HF_TOKEN)
 #     return response
 
 
-def predict_image(image_url, text):
+def extract_image_from_pdf(pdf_url, dpi=75):
+    """
+    Extract first page of PDF as image in memory
+    
+    Args:
+        pdf_url (str): URL of PDF
+        dpi (int): Image resolution
+    
+    Returns:
+        PIL.Image: First page as image or None
+    """
+    try:
+        # Download PDF
+        response = requests.get(pdf_url, timeout=30)
+        response.raise_for_status()
+        
+        # Open PDF from bytes
+        pdf_document = fitz.open(stream=response.content, filetype="pdf")
+        
+        # Get first page
+        first_page = pdf_document[0]
+        
+        # Render page to pixmap
+        pix = first_page.get_pixmap(matrix=fitz.Matrix(dpi/72, dpi/72))
+        
+        # Convert to PIL Image
+        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+        
+        pdf_document.close()
+        return img
+    
+    except Exception as e:
+        print(f"Error extracting first page: {e}")
+        return None
+
+
+
+def predict_image(image_url, text, file_pref):
     try:
         # Download the image from the URL
         # response = requests.get(image_url)
         # response.raise_for_status()  # Raise an error for invalid responses
         # image = Image.open(io.BytesIO(response.content)).convert("RGB")
-
-        response = requests.get(image_url)
-        response.raise_for_status()  # Raise an error for invalid responses
-        image = Image.open(io.BytesIO(response.content)).convert("RGB")
-        
+        if file_pref == 'img':
+            response = requests.get(image_url)
+            response.raise_for_status()  # Raise an error for invalid responses
+            image = Image.open(io.BytesIO(response.content)).convert("RGB")
+        else: 
+            image = extract_image_from_pdf(image_url)
+            
         messages = [
             {"role": "user", "content": [
                 {"type": "image"},  # Specify that an image is provided
@@ -200,6 +240,7 @@ PROMPT = (
     "'Course Name:'\n"
     "'Course Description:'\n"
     "'Course Credits:'\n"
+    "'Course Learning Outcomes:'\n"
     "'Delivery Method:'\n"
     "'Prerequisite(s):'\n"
     "'Co-requisite(s):'\n"
@@ -263,7 +304,8 @@ def extract_info():
         if data["img_url"] is not None:
             prompt_skills = f"{PROMPT_IMAGE}\n"
             img_url = data["img_url"]
-            response_image = predict_image(img_url, prompt_skills)
+            file_pref = data["file_pref"]
+            response_image = predict_image(img_url, prompt_skills, file_pref)
         else:
             response_image = ''