Update main.py

main.py

@@ -15,16 +15,26 @@ from io import BytesIO
 from PIL import Image
 import io
 import requests
+import fitz  # PyMuPDF
+import os
+
 
 from dotenv import load_dotenv
 # Load the environment variables from the .env file
 load_dotenv()
 
-secret = os.environ["GEMINI"]
-genai.configure(api_key=secret)
-model_vision = genai.GenerativeModel('gemini-pro-vision')
+# Configure Gemini API
+genai.configure(api_key="AIzaSyBsutShR1tNNdomkaL3DYHjMrM_59Y1mg8")
+#secret = os.environ["GEMINI"]
+#genai.configure(api_key=secret)
+model_vision = genai.GenerativeModel('gemini-1.5-flash')
 model_text = genai.GenerativeModel('gemini-pro')
 
+
+
+
+
+
 app = FastAPI()
 
 app.add_middleware(
@@ -37,56 +47,35 @@ app.add_middleware(
 
 
 
-def encode_image(image):
-    # Convert image to BytesIO object (in memory)
-    buffered = BytesIO()
-    image.save(buffered, format=image.format)  # Use the original image format (e.g., PNG, JPEG)
-    img_bytes = buffered.getvalue()
-
-    # Encode image to base64
-    base64_image = base64.b64encode(img_bytes).decode('utf-8')
-    return base64_image
 
 
-
-def vision(image):
-    # OpenAI API Key
-    api_key = os.environ["OPEN_AI"]
-    
-
-    # Getting the base64 string
-    base64_image = encode_image(image)
-    
-    headers = {
-      "Content-Type": "application/json",
-      "Authorization": f"Bearer {api_key}"
-    }
-    
-    payload = {
-      "model": "gpt-4o-mini",
-      "messages": [
-        {
-          "role": "user",
-          "content": [
-            {
-              "type": "text",
-              "text": "extract all data from this image"
-            },
-            {
-              "type": "image_url",
-              "image_url": {
-                "url": f"data:image/jpeg;base64,{base64_image}"
-              }
-            }
-          ]
-        }
-      ],
-      "max_tokens": 300
-    }
+def vision(file_content):
+    # Open the PDF
+    pdf_document = fitz.open("pdf",pdf)
+    gemini_input = ["extract the whole text"]
+    # Iterate through the pages
+    for page_num in range(len(pdf_document)):
+        # Select the page
+        page = pdf_document.load_page(page_num)
+        
+        # Render the page to a pixmap (image)
+        pix = page.get_pixmap()
+        print(type(pix))
+        
+        # Convert the pixmap to bytes
+        img_bytes = pix.tobytes("png")
+        
+        # Convert bytes to a PIL Image
+        img = Image.open(io.BytesIO(img_bytes))
+        gemini_input.append(img)
+        # # Save the image if needed
+        # img.save(f'page_{page_num + 1}.png')
     
-    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+    print("PDF pages converted to images successfully!")
     
-    return response.json()['choices'][0]['message']['content']
+    # Now you can pass the PIL image to the model_vision
+    response = model_vision.generate_content(gemini_input).text
+    return response
 
 
 @app.post("/get_ocr_data/")
@@ -99,16 +88,13 @@ async def get_data(input_file: UploadFile = File(...)):
         text = ""
 
         if file_type == "application/pdf":
-            # Read PDF file using PyPDF2
-            pdf_reader = PdfReader(io.BytesIO(file_content))
-            for page in pdf_reader.pages:
-                text += page.extract_text()
-        
-        elif file_type in ["image/jpeg", "image/png", "image/jpg"]:
-            # Read Image file using PIL and pytesseract
-            image = Image.open(io.BytesIO(file_content))
-            text = vision(image)
-        
+            if text=="":
+               text = vision(file_content)
+            # else:
+            #     # Read PDF file using PyPDF2
+            #     pdf_reader = PdfReader(io.BytesIO(file_content))
+            #     for page in pdf_reader.pages:
+            #         text += page.extract_text()
         else:
             raise HTTPException(status_code=400, detail="Unsupported file type")