Update

app.py

@@ -1,2 +1,61 @@
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+import streamlit as st
+from layout_extractor import convert_pdf_to_images, analyze_layout, extract_text_from_blocks, extract_key_values
+from processor import load_images, analyze_layout, extract_text_from_blocks, rule_based_kv_extraction
+import json
+
+st.set_page_config(page_title="Document AI", layout="wide")
+st.title("🧠 AI-Driven Document Layout & Info Extractor")
+
+uploaded_file = st.file_uploader("Upload a PDF or Image", type=["pdf", "png", "jpg", "jpeg"])
+
+if uploaded_file:
+    images = load_images(uploaded_file)
+    for i, image in enumerate(images):
+        st.subheader(f"Page {i+1}")
+        st.image(image, use_column_width=True)
+
+        with st.spinner("Analyzing layout..."):
+            layout = analyze_layout(image)
+            blocks = extract_text_from_blocks(image, layout)
+            kv_data = rule_based_kv_extraction(blocks)
+
+        st.success("Done! Here's what we found:")
+        st.json(kv_data)
+
+        st.subheader("✏️ Edit Extracted Fields")
+        edited_data = {}
+        for key, value in kv_data.items():
+            edited_data[key] = st.text_input(f"{key}", value)
+
+        st.download_button("⬇️ Download JSON", data=json.dumps(edited_data, indent=2),
+                           file_name="extracted_data.json", mime="application/json")
+
+        with st.expander("🔍 All Detected Segments"):
+            for b in blocks:
+                st.markdown(f"**{b['type']}**: {b['text'][:150]}...")
+
+st.title("📄 AI-Driven Document Layout Analyzer")
+
+uploaded_file = st.file_uploader("Upload a PDF or Image", type=["pdf", "png", "jpg", "jpeg"])
+
+if uploaded_file:
+    if uploaded_file.name.endswith(".pdf"):
+        images = convert_pdf_to_images(uploaded_file)
+    else:
+        from PIL import Image
+        images = [Image.open(uploaded_file)]
+
+    for i, image in enumerate(images):
+        st.image(image, caption=f"Page {i+1}", use_column_width=True)
+        layout = analyze_layout(image)
+        blocks = extract_text_from_blocks(image, layout)
+        key_values = extract_key_values(blocks)
+
+        st.subheader("Extracted Key Data")
+        st.json(key_values)
+
+        st.subheader("All Segments")
+        for block in blocks:
+            st.markdown(f"**{block['type']}**: {block['text'][:200]}...")
+
+        st.download_button("Download JSON", data=json.dumps(key_values, indent=2), file_name="extracted_data.json")

layout_extractor.py

@@ -0,0 +1,41 @@
+import layoutparser as lp
+import pytesseract
+import json
+from pdf2image import convert_from_path
+from PIL import Image
+
+def convert_pdf_to_images(pdf_path):
+    return convert_from_path(pdf_path)
+
+def analyze_layout(image):
+    model = lp.Detectron2LayoutModel(
+        config_path='lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config',
+        extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8],
+        label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}
+    )
+    layout = model.detect(image)
+    return layout
+
+def extract_text_from_blocks(image, layout):
+    blocks = []
+    for block in layout:
+        segment_image = block.crop_image(image)
+        text = pytesseract.image_to_string(segment_image)
+        blocks.append({
+            "type": block.type,
+            "text": text.strip(),
+            "coordinates": block.coordinates
+        })
+    return blocks
+
+def extract_key_values(blocks):
+    data = {}
+    for block in blocks:
+        text = block["text"]
+        if "invoice" in text.lower():
+            data["invoice_number"] = text
+        elif "total" in text.lower():
+            data["total_amount"] = text
+        elif "customer" in text.lower():
+            data["customer_name"] = text
+    return data

processor.py

@@ -0,0 +1,45 @@
+import layoutparser as lp
+import pytesseract
+from pdf2image import convert_from_path
+from PIL import Image
+import json
+
+model = lp.Detectron2LayoutModel(
+    config_path="lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config",
+    extra_config=["MODEL.ROI_HEADS.SCORE_THRESH_TEST", 0.8],
+    label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
+)
+
+def load_images(uploaded_file):
+    if uploaded_file.name.endswith(".pdf"):
+        return convert_from_path(uploaded_file)
+    else:
+        return [Image.open(uploaded_file)]
+
+def analyze_layout(image):
+    layout = model.detect(image)
+    return layout
+
+def extract_text_from_blocks(image, layout):
+    blocks = []
+    for block in layout:
+        cropped = block.crop_image(image)
+        text = pytesseract.image_to_string(cropped)
+        blocks.append({
+            "type": block.type,
+            "text": text.strip(),
+            "coords": block.coordinates
+        })
+    return blocks
+
+def rule_based_kv_extraction(blocks):
+    data = {}
+    for b in blocks:
+        t = b["text"].lower()
+        if "invoice" in t and "number" in t:
+            data["Invoice Number"] = b["text"]
+        elif "total" in t:
+            data["Total Amount"] = b["text"]
+        elif "customer" in t:
+            data["Customer Name"] = b["text"]
+    return data

requirements.txt

@@ -0,0 +1,8 @@
+streamlit
+layoutparser
+pdf2image
+pytesseract
+transformers
+torch
+Pillow
+opencv-python