import streamlit as st import base64 import requests from PIL import Image, ImageDraw from io import BytesIO import fitz # PyMuPDF import time # Configuration - Get API key from Streamlit secrets GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"] GEMINI_MODEL = "gemini-2.0-flash" DOCUMENT_TYPES = ["Land Records", "Caste Certificates", "Property Registrations", "Others"] # Initialize session state (excluding widget-controlled keys) def initialize_session_state(): if "chat_history" not in st.session_state: st.session_state["chat_history"] = [] if "processed_doc" not in st.session_state: st.session_state["processed_doc"] = None if "doc_preview" not in st.session_state: st.session_state["doc_preview"] = None # Reset session state (excluding widget-controlled keys) def reset_session_state(): for key in ["chat_history", "processed_doc", "doc_preview"]: st.session_state.pop(key, None) # Encode uploaded file to base64 def encode_file(uploaded_file): try: file_bytes = uploaded_file.getvalue() if uploaded_file.type == "application/pdf": pdf = fitz.open(stream=BytesIO(file_bytes)) page = pdf[0] pix = page.get_pixmap() img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) elif uploaded_file.type.startswith("image/"): img = Image.open(BytesIO(file_bytes)) elif uploaded_file.type == "text/plain": text = file_bytes.decode("utf-8") img = Image.new("RGB", (800, 600), color=(73, 109, 137)) d = ImageDraw.Draw(img) d.text((10, 10), text, fill=(255, 255, 0)) else: st.error("Unsupported file format") return None img_byte_arr = BytesIO() img.save(img_byte_arr, format="JPEG") return base64.b64encode(img_byte_arr.getvalue()).decode("utf-8") except Exception as e: st.error(f"File processing error: {str(e)}") return None # Query Gemini API def query_gemini(prompt, image_b64=None): url = f"https://generativelanguage.googleapis.com/v1/models/{GEMINI_MODEL}:generateContent?key={GEMINI_API_KEY}" parts = [{"text": prompt}] if image_b64: parts.append({ "inline_data": { "mime_type": "image/jpeg", "data": image_b64 } }) try: response = requests.post( url, json={"contents": [{"parts": parts}]}, headers={"Content-Type": "application/json"}, timeout=30 ) if response.status_code != 200: st.error(f"API Request failed with status code: {response.status_code}") return None data = response.json() if "error" in data: st.error(f"API Error: {data['error'].get('message', 'Unknown error')}") return None if not data.get("candidates"): st.error("No response candidates found in API response") return None candidate = data["candidates"][0] return candidate.get("content", {}).get("parts", [{}])[0].get("text", "No response text found") except requests.exceptions.RequestException as e: st.error(f"API Request failed: {str(e)}") return None except Exception as e: st.error(f"Unexpected error: {str(e)}") return None # Process the uploaded document def process_document(): uploaded_file = st.session_state.get("uploaded_file") if not uploaded_file: st.error("Please upload a document first.") return try: with st.spinner("Analyzing document..."): # Encode file to base64 image_b64 = encode_file(uploaded_file) if not image_b64: return # Store preview image if uploaded_file.type == "application/pdf": pdf = fitz.open(stream=BytesIO(uploaded_file.getvalue())) page = pdf[0] pix = page.get_pixmap() st.session_state.doc_preview = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) else: st.session_state.doc_preview = Image.open(uploaded_file) # Classify document classify_prompt = ( f"Classify this document into one of these categories: {', '.join(DOCUMENT_TYPES)}. " "Respond only with the category name." ) doc_type = query_gemini(classify_prompt, image_b64) # Extract details extract_prompt = ( "Extract key details including:\n" "- Names\n" "- Dates\n" "- Identification numbers\n" "- Locations\n" "Format as a bullet-point list." ) details = query_gemini(extract_prompt, image_b64) # Verify authenticity verify_prompt = "Analyze this document for signs of tampering. Provide verification status in short(2 Lines)." verification = query_gemini(verify_prompt, image_b64) st.session_state.processed_doc = { "type": doc_type or "Unclassified", "details": details or "No details extracted", "verification": verification or "Verification failed", } st.success("Document processing complete!") time.sleep(1) except Exception as e: st.error(f"Document processing failed: {str(e)}") st.session_state.processed_doc = None # Main app function def main(): st.set_page_config(page_title="DocVerify AI", layout="wide") initialize_session_state() # Sidebar Controls with st.sidebar: st.header("Document Controls") # The file uploader widget manages its own state with key "uploaded_file" st.file_uploader( "Upload Document", type=["pdf", "jpg", "jpeg", "png", "txt"], key="uploaded_file", on_change=process_document, help="Supported formats: PDF, JPG, PNG, TXT" ) if st.button("New Document"): reset_session_state() st.experimental_rerun() if st.session_state.get("processed_doc"): st.divider() st.subheader("Document Summary") st.markdown(f"**Type:** {st.session_state.processed_doc['type']}") st.markdown(f"**Verification Status:** {st.session_state.processed_doc['verification']}") # Main Interface st.title("📄 Automated Document Verifier") if st.session_state.get("processed_doc") and st.session_state.get("doc_preview"): col1, col2 = st.columns([1, 2]) with col1: st.subheader("Document Preview") st.image(st.session_state.doc_preview, use_column_width=True) with col2: st.subheader("Extracted Details") st.markdown(st.session_state.processed_doc["details"]) st.subheader("Verification Analysis") st.markdown(st.session_state.processed_doc["verification"]) else: st.info("Please upload a document to begin verification analysis") if __name__ == "__main__": main()