Update app.py

app.py

@@ -11,7 +11,7 @@ GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
 GEMINI_MODEL = "gemini-2-flash"
 DOCUMENT_TYPES = ["Land Records", "Caste Certificates", "Property Registrations"]
 
-# Initialize session state
+# Initialize session state (excluding widget-controlled keys)
 def initialize_session_state():
     if "chat_history" not in st.session_state:
         st.session_state["chat_history"] = []
@@ -19,12 +19,10 @@ def initialize_session_state():
         st.session_state["processed_doc"] = None
     if "doc_preview" not in st.session_state:
         st.session_state["doc_preview"] = None
-    if "uploaded_file" not in st.session_state:
-        st.session_state["uploaded_file"] = None
 
-# Reset session state
+# Reset session state (excluding widget-controlled keys)
 def reset_session_state():
-    for key in ["chat_history", "processed_doc", "doc_preview", "uploaded_file"]:
+    for key in ["chat_history", "processed_doc", "doc_preview"]:
         st.session_state.pop(key, None)
 
 # Encode uploaded file to base64
@@ -37,11 +35,11 @@ def encode_file(uploaded_file):
             page = pdf[0]
             pix = page.get_pixmap()
             img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        elif uploaded_file.type.startswith('image/'):
+        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))
+            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:
@@ -49,8 +47,9 @@ def encode_file(uploaded_file):
             return None
 
         img_byte_arr = BytesIO()
-        img.save(img_byte_arr, format='JPEG')
-        return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+        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
@@ -58,9 +57,15 @@ def encode_file(uploaded_file):
 # 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}})
+        parts.append({
+            "inline_data": {
+                "mime_type": "image/jpeg",
+                "data": image_b64
+            }
+        })
     
     try:
         response = requests.post(
@@ -74,16 +79,17 @@ def query_gemini(prompt, image_b64=None):
             return None
         
         data = response.json()
-        if 'error' in data:
+        if "error" in data:
             st.error(f"API Error: {data['error'].get('message', 'Unknown error')}")
             return None
-        
-        if not data.get('candidates'):
+            
+        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')
+        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
@@ -93,26 +99,46 @@ def query_gemini(prompt, image_b64=None):
 
 # Process the uploaded document
 def process_document():
-    if not st.session_state.uploaded_file:
+    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..."):
-            image_b64 = encode_file(st.session_state.uploaded_file)
+            # Encode file to base64
+            image_b64 = encode_file(uploaded_file)
             if not image_b64:
                 return
 
-            classify_prompt = f"Classify this document into one of these categories: {', '.join(DOCUMENT_TYPES)}. Respond only with the category name."
+            # 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_prompt = """Extract key details including:
-            - Names
-            - Dates
-            - Identification numbers
-            - Locations
-            Format as a bullet-point list."""
+            # 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."
             verification = query_gemini(verify_prompt, image_b64)
 
@@ -121,8 +147,10 @@ def process_document():
                 "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
@@ -132,20 +160,45 @@ def main():
     st.set_page_config(page_title="DocVerify AI", layout="wide")
     initialize_session_state()
     
-    st.sidebar.header("Document Controls")
-    st.sidebar.file_uploader("Upload Document", type=["pdf", "jpg", "jpeg", "png", "txt"], key="uploaded_file", on_change=process_document)
-    
-    if st.sidebar.button("New Document"):
-        reset_session_state()
-        st.rerun()
-    
-    st.title("DocVerify AI - Document Analysis")
+    # 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.processed_doc:
-        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']}")
-        st.text_area("Extracted Details", st.session_state.processed_doc['details'], height=200)
+    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()