Spaces:
Running
Running
File size: 6,830 Bytes
7b76e52 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 |
import streamlit as st
import base64
import requests
from PIL import Image, ImageDraw
from io import BytesIO
import fitz # PyMuPDF
import time
import os
# Configuration
GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
GEMINI_MODEL = "gemini-1.5-flash" # Verify with Gemini API documentation
DOCUMENT_TYPES = ["Land Records", "Caste Certificates", "Property Registrations"]
# Define a class to mimic Streamlit's UploadedFile object for the example image
class ExampleFile:
def __init__(self, file_path):
self.name = os.path.basename(file_path)
self.type = "image/jpeg"
with open(file_path, "rb") as f:
self._buffer = BytesIO(f.read())
def getvalue(self):
return self._buffer.getvalue()
# Initialize session state
def initialize_session_state():
if "processed_doc" not in st.session_state:
st.session_state.processed_doc = None
if "current_file" not in st.session_state:
st.session_state.current_file = None
# Reset session state
def reset_session_state():
st.session_state.processed_doc = None
st.session_state.current_file = None
# Encode file to base64 for Gemini API
def encode_file(file):
try:
file_content = file.getvalue()
return base64.b64encode(file_content).decode('utf-8')
except Exception as e:
st.error(f"Error encoding file: {str(e)}")
return None
# Query Gemini API
def query_gemini(prompt, image_b64):
try:
headers = {
"Authorization": f"Bearer {GEMINI_API_KEY}",
"Content-Type": "application/json"
}
payload = {
"model": GEMINI_MODEL,
"prompt": prompt,
"image": image_b64
}
response = requests.post(
"https://api.gemini.com/v1/analyze", # Replace with actual Gemini API endpoint
headers=headers,
json=payload
)
response.raise_for_status()
return response.json().get("result", "")
except Exception as e:
st.error(f"Gemini API error: {str(e)}")
return None
# Process the document
def process_document(file):
try:
with st.spinner("Analyzing document..."):
# Encode file to base64
image_b64 = encode_file(file)
if not image_b64:
return
# Store preview image
if file.type == "application/pdf":
pdf = fitz.open(stream=BytesIO(file.getvalue()))
page = pdf[0]
pix = page.get_pixmap()
st.session_state.doc_preview = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
elif file.type.startswith('image/'):
st.session_state.doc_preview = Image.open(BytesIO(file.getvalue()))
elif file.type == "text/plain":
text = file.getvalue().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))
st.session_state.doc_preview = img
else:
st.error("Unsupported file format")
return
# 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 and organize all important details from this document including:
- Names
- Dates
- Identification numbers
- Locations
- Key terms
Format as a bullet-point list with clear headings."""
details = query_gemini(extract_prompt, image_b64)
# Verify authenticity
verify_prompt = "Analyze this document for signs of tampering or forgery. Check for: inconsistent fonts, alignment issues, suspicious modifications. Provide verification conclusion."
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",
"preview": st.session_state.doc_preview
}
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 application
def main():
st.set_page_config(page_title="DocVerify AI", layout="wide")
initialize_session_state()
# Sidebar Controls
with st.sidebar:
st.header("Document Controls")
uploaded_file = st.file_uploader(
"Upload Document",
type=["pdf", "jpg", "jpeg", "png", "txt"],
key="uploaded_file"
)
if st.button("Use Example Image"):
example_file_path = "Caste-Certificate.jpg"
st.session_state.current_file = ExampleFile(example_file_path)
elif uploaded_file:
st.session_state.current_file = uploaded_file
else:
st.session_state.current_file = None
if st.button("Process Document"):
if st.session_state.current_file:
process_document(st.session_state.current_file)
else:
st.error("Please select a document to process.")
if st.button("New Document"):
reset_session_state()
st.rerun()
if st.session_state.processed_doc:
st.divider()
st.subheader("Document Summary")
st.markdown(f"**Type:** {st.session_state.processed_doc['type']}")
st.markdown(f"**Verification Status:**\n{st.session_state.processed_doc['verification']}")
# Main Interface
st.title("📄 Automated Document Verifier")
if st.session_state.processed_doc and 'preview' in st.session_state.processed_doc:
col1, col2 = st.columns([1, 2])
with col1:
st.subheader("Document Preview")
st.image(st.session_state.processed_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 select a document and click 'Process Document' to begin verification analysis.")
if __name__ == "__main__":
main() |