Igniteit/ocrdeploytest

1try: from pip._internal.operations import freeze2except ImportError: # pip < 10.03    from pip.operations import freeze4 5pkgs = freeze.freeze()6for pkg in pkgs: print(pkg)7import os 8from fastapi import FastAPI, HTTPException, File, UploadFile,Query9from fastapi.middleware.cors import CORSMiddleware10from PyPDF2 import PdfReader11import google.generativeai as genai12import json13from PIL import Image14import io15import requests16import fitz  # PyMuPDF17import os18 19 20from dotenv import load_dotenv21# Load the environment variables from the .env file22load_dotenv()23 24# Configure Gemini API25secret = os.environ["GEMINI"]26genai.configure(api_key=secret)27model_vision = genai.GenerativeModel('gemini-1.5-flash')28model_text = genai.GenerativeModel('gemini-pro')29 30 31 32 33 34 35app = FastAPI()36 37app.add_middleware(38    CORSMiddleware,39    allow_origins=["*"],40    allow_credentials=True,41    allow_methods=["*"],42    allow_headers=["*"],43)44 45 46 47 48 49def vision(file_content):50    # Open the PDF51    pdf_document = fitz.open("pdf",file_content)52    gemini_input = ["extract the whole text"]53    # Iterate through the pages54    for page_num in range(len(pdf_document)):55        # Select the page56        page = pdf_document.load_page(page_num)57        58        # Render the page to a pixmap (image)59        pix = page.get_pixmap()60        print(type(pix))61        62        # Convert the pixmap to bytes63        img_bytes = pix.tobytes("png")64        65        # Convert bytes to a PIL Image66        img = Image.open(io.BytesIO(img_bytes))67        gemini_input.append(img)68        # # Save the image if needed69        # img.save(f'page_{page_num + 1}.png')70    71    print("PDF pages converted to images successfully!")72    73    # Now you can pass the PIL image to the model_vision74    response = model_vision.generate_content(gemini_input).text75    return response76 77 78@app.post("/get_ocr_data/")79async def get_data(user_id: str = Query(...),input_file: UploadFile = File(...)):80    #try:81        # Determine the file type by reading the first few bytes82        file_content = await input_file.read()83        file_type = input_file.content_type84        85        text = ""86 87        if file_type == "application/pdf":88                # Read PDF file using PyPDF289                pdf_reader = PdfReader(io.BytesIO(file_content))90                for page in pdf_reader.pages:91                    text += page.extract_text()92                    93                if len(text)<10:94                   print("vision called")95                   text = vision(file_content)96        else:97            raise HTTPException(status_code=400, detail="Unsupported file type")98        99        100 101        # Call Gemini (or another model) to extract required data102        prompt = f"""This is CV data: {text.strip()} 103                IMPORTANT: The output should be a JSON array! Make Sure the JSON is valid.104                                                                  105                Example Output:106                [107                    "firstname" : "firstname",108                    "lastname" : "lastname",109                    "email" : "email",110                    "contact_number" : "contact number",111                    "home_address" : "full home address",112                    "home_town" : "home town or city",113                    "total_years_of_experience" : "total years of experience",114                    "education": "Institution Name, Degree Name,115                    "LinkedIn_link" : "LinkedIn link",116                    "positions": [ "Job title 1", "Job title 2", "Job title 3" ],117                    "industry": "[ "industry 1", "industry 2", "industry 3" ],  # List all industries the candidate has worked in, inferred from job titles, companies, or experience",118                    "experience" : "experience",119                    "skills" : skills(Identify and list specific skills mentioned in both the skills section and inferred from the experience section)120                ]121                """122        123        response = model_text.generate_content(prompt)124        print(response.text)125        data = json.loads(response.text.replace("JSON", "").replace("json", "").replace("```", ""))126        return {"data": data}127 128    #except Exception as e:129        #raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")