OrbitGuy2244/RAG_PDF

1import streamlit as st2import json3# import pytesseract4from google.cloud import vision5import fitz  # PyMuPDF for PDFs6from PIL import Image7import google.generativeai as genai8import pandas as pd9from dotenv import load_dotenv10import os11 12load_dotenv()13# GOOGLE_API_KEY = os.getenv("Google_API")14GOOGLE_API_KEY = os.environ.get("Google_API")15 16# Configure API Key17genai.configure(api_key=GOOGLE_API_KEY)18 19# Set Google Cloud credentials path20os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "gen-lang-client-0757522976-213a1e391699.json"21 22# Load the Gemini model23gemini_model = genai.GenerativeModel("gemini-2.0-flash-exp")24 25def extract_text_from_pdf(pdf_file):26    """Extracts text from PDF with better accuracy."""27    doc = fitz.open(stream=pdf_file.read(), filetype="pdf")28    text = "\n".join([page.get_text("text") for page in doc])29    return text.strip()30 31# def extract_text_from_image(image_file):32#     """Extracts text from an image using Tesseract OCR with optimized settings."""33#     image = Image.open(image_file)34#     text = pytesseract.image_to_string(image, config="--oem 3 --psm 6")35#     return text.strip()36def extract_text_from_image(image_file):37    """Extracts text using Google Cloud Vision API."""38    client = vision.ImageAnnotatorClient()39    content = image_file.read()40    image = vision.Image(content=content)41    response = client.text_detection(image=image)42    text = response.full_text_annotation.text43    return text.strip()44 45def clean_extracted_text(text):46    """Cleans extracted text to remove unwanted artifacts."""47    text = text.replace("\n\n", "\n").strip()  # Remove extra new lines48    return text49 50def get_json_from_text(text):51    """Uses Google Gemini API to structure extracted text into JSON format."""52    prompt = f"""53    The following text is extracted from a structured form. Convert it into a properly formatted JSON with relevant key-value pairs.54 55    Text:56    {text}57 58    - Extract all key-value pairs correctly.59    - Ensure fields like "Name", "Date of Birth", "Account Number", etc., are well-organized.60    - Return **only** JSON format without additional text.61    """62 63    response = gemini_model.generate_content(prompt)64 65    if response and hasattr(response, "text"):66        response_text = response.text.strip()67 68        # Handle cases where Gemini wraps JSON in markdown code blocks69        if response_text.startswith("```json"):70            response_text = response_text.replace("```json", "").replace("```", "").strip()71 72        try:73            json_data = json.loads(response_text)  # Parse JSON safely74        except json.JSONDecodeError:75            json_data = {"error": "Failed to parse response as JSON. Please check the extraction."}76    else:77        json_data = {"error": "No response from Google Gemini API."}78 79    return json_data80 81def extract_medical_terms(text):82    """Uses Google Gemini API to extract and categorize medical terms."""83    prompt = f"""84    Extract only medical-related terms from the following text. Examples include diseases, symptoms, treatments, and medical conditions85    and categorize all medical-related terms from the following text.86    The categories should include:87    - Diseases88    - Medications89    - Symptoms90    - Medical Procedures91    - Medical Devices92    - Chemicals93 94    Text:95    {text}96    97    Return the response in this **exact JSON format**:98    {{99        "Diseases": ["disease1", "disease2"],100        "Medications": ["medication1", "medication2"],101        "Symptoms": ["symptom1", "symptom2"],102        "Medical Procedures": ["procedure1", "procedure2"],103        "Medical Devices": ["device1", "device2"],104        "Chemicals": ["chemical1", "chemical2"]105    }}106 107    Ensure:108    - No extra text, only valid JSON.109    - No explanations, only the structured JSON response.110    """111 112    response = gemini_model.generate_content(prompt)113 114    if response and hasattr(response, "text"):115        response_text = response.text.strip()116 117        # Handle cases where Gemini wraps JSON in markdown code blocks118        if response_text.startswith("```json"):119            response_text = response_text.replace("```json", "").replace("```", "").strip()120 121        try:122            json_data = json.loads(response_text)  # Parse JSON safely123            return json_data124        except json.JSONDecodeError:125            return {"error": "Gemini returned an invalid response. Try again."}126    else:127        return {"error": "No response from Google Gemini API."}128 129# Streamlit UI Enhancements130st.set_page_config(page_title="AI Form Extractor", layout="centered")131st.markdown("""132    <style>133    .stButton > button {134        background-color: #590f79;135        color: white;136        font-size: 16px;137        border-radius: 10px;138        padding: 10px;139        border: none;140    }141    .stButton > button:hover {142        background-color: #590f79;143    }144    .stTextArea > label {145        font-size: 18px;146        font-weight: bold;147    }148    .stJson > label {149        font-size: 18px;150        font-weight: bold;151    }152    .center-text {153        text-align: center;154        font-size: 18px;155        font-weight: bold;156        color: #4CAF50;157    }158    </style>159    """, unsafe_allow_html=True)160 161st.title("📄 RAG Based Form Extraction to JSON")162st.write("Upload a **PDF or image** of a form, and the system will extract and structure the details into JSON.")163 164# Initialize session state for UI buttons165if "show_extracted" not in st.session_state:166    st.session_state.show_extracted = False167if "show_medical_terms" not in st.session_state:168    st.session_state.show_medical_terms = False169 170uploaded_file = st.file_uploader("📂 Upload a form (PDF or Image)", type=["pdf", "png", "jpg", "jpeg"], help="Supported formats: PDF, PNG, JPG, JPEG")171 172if uploaded_file is not None:173    file_type = uploaded_file.type174    extracted_text = ""175 176    with st.spinner("🔍 Extracting text..."):177        if "pdf" in file_type:178            extracted_text = extract_text_from_pdf(uploaded_file)179        else:180            extracted_text = extract_text_from_image(uploaded_file)181 182    cleaned_text = clean_extracted_text(extracted_text)183 184    # Button to show extracted text185    if st.button("👁 Show Extracted Text"):186        st.session_state.show_extracted = not st.session_state.show_extracted187 188    if st.session_state.show_extracted:189        st.subheader("📜 Extracted Text")190        st.text_area("", cleaned_text, height=200)191 192    # Button to convert text to JSON193    if st.button("🛠 Convert to JSON"):194        with st.spinner("🧠 Generating structured JSON..."):195            json_output = get_json_from_text(cleaned_text)196        st.subheader("📊 Structured JSON Output")197        st.json(json_output)198        199        # Add download button for JSON file200        json_str = json.dumps(json_output, indent=4)201        st.download_button(label="📥 Download JSON", data=json_str, file_name="extracted_data.json", mime="application/json")202 203    # Button to extract medical terms204    if st.button("⚕️ Extract Medical Terms"):205        st.session_state.show_medical_terms = not st.session_state.show_medical_terms206 207    if st.session_state.show_medical_terms:208        with st.spinner("🔎 Identifying Medical Terms..."):209            medical_terms = extract_medical_terms(cleaned_text)210        st.subheader("🩺 Medical Terms Found")211        st.write(medical_terms if medical_terms else "No medical terms found.")212