Update main.py

main.py

@@ -8,10 +8,12 @@ import os
 from fastapi import FastAPI, HTTPException, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from PyPDF2 import PdfReader
-
-
 import google.generativeai as genai
 import json
+import base64
+from io import BytesIO
+from PIL import Image'
+import requests
 
 secret = os.environ["key"]
 genai.configure(api_key=secret)
@@ -28,30 +30,95 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+
+
+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 = "sk-proj-1j1aFDCU8KrWAeFMAGPPT3BlbkFJ6rDxGgu8C99E3Wh6siUs"
+    
+
+    # 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
+    }
+    
+    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+    
+    print(response.json()['choices'][0]['message']['content'])
+
+
 @app.post("/get_ocr_data/")
-async def get_data(pdf: UploadFile = File(...)):
+async def get_data(input_file: UploadFile = File(...)):
     try:
-        # Read PDF file using PyPDF2
-        pdf_reader = PdfReader(pdf.file)
+        # Determine the file type by reading the first few bytes
+        file_content = await input_file.read()
+        file_type = input_file.content_type
+        
         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))
+            return encode_image(image)
+            text = vision(image)
         
-        # Extract text from each page
-        for page in pdf_reader.pages:
-            text += page.extract_text()
+        else:
+            raise HTTPException(status_code=400, detail="Unsupported file type")
 
-        # call gemini to get required data extracted text
-        prompt = f"""this is cv data : {text.strip()} 
-         i want only 
+        # Call Gemini (or another model) to extract required data
+        prompt = f"""This is CV data: {text.strip()} 
+         I want only:
 
-         fisrtname,lastname,contact number,total years of experince,linkdn link,experinece,skils 
+         firstname, lastname, contact number, total years of experience, LinkedIn link, experience, skills
 
-         in json format only"""
+         in JSON format only"""
         
         response = model_text.generate_content(prompt)
-        data = json.loads(response.text.replace("```json","").replace("```",""))
-        return {"data":data}
+        data = json.loads(response.text.replace("```json", "").replace("```", ""))
+        return {"data": data}
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing PDF: {str(e)}")
-
-
+        raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")