Update app.py

app.py

@@ -1,140 +1,169 @@
 # app.py
 
 import streamlit as st
-import numpy as np
 import cv2
+import numpy as np
 from PIL import Image, ImageEnhance
-from fpdf import FPDF
 import io
-import base64
-import os
-
-# Function to detect and crop the document
-def detect_and_crop(img):
-    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    blur = cv2.GaussianBlur(gray, (5,5), 0)
-    edged = cv2.Canny(blur, 50, 200)
-
-    contours, _ = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if len(contours) == 0:
-        return img  # Return original if nothing found
-
-    c = max(contours, key=cv2.contourArea)
-    peri = cv2.arcLength(c, True)
-    approx = cv2.approxPolyDP(c, 0.02 * peri, True)
-
-    if len(approx) == 4:
-        pts = approx.reshape(4, 2)
-        rect = order_points(pts)
-        (tl, tr, br, bl) = rect
-
-        widthA = np.linalg.norm(br - bl)
-        widthB = np.linalg.norm(tr - tl)
-        maxWidth = max(int(widthA), int(widthB))
-
-        heightA = np.linalg.norm(tr - br)
-        heightB = np.linalg.norm(tl - bl)
-        maxHeight = max(int(heightA), int(heightB))
-
-        dst = np.array([
-            [0, 0],
-            [maxWidth - 1, 0],
-            [maxWidth - 1, maxHeight - 1],
-            [0, maxHeight - 1]
-        ], dtype="float32")
-
-        M = cv2.getPerspectiveTransform(rect, dst)
-        warped = cv2.warpPerspective(img, M, (maxWidth, maxHeight))
-
-        return warped
+from fpdf import FPDF
+import datetime
+
+# -----------------------
+# Helper Functions
+# -----------------------
+
+def auto_crop_document(image):
+    """Auto-detect document and crop it."""
+    img = np.array(image)
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    edged = cv2.Canny(blurred, 75, 200)
+
+    contours, _ = cv2.findContours(edged.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:5]
+
+    for c in contours:
+        peri = cv2.arcLength(c, True)
+        approx = cv2.approxPolyDP(c, 0.02 * peri, True)
+
+        if len(approx) == 4:
+            doc_cnt = approx
+            break
     else:
-        return img
+        return image  # Return original if no contour found
 
-# Helper to order contour points
-def order_points(pts):
+    pts = doc_cnt.reshape(4, 2)
     rect = np.zeros((4, 2), dtype="float32")
+
     s = pts.sum(axis=1)
-    rect[0] = pts[np.argmin(s)]
-    rect[2] = pts[np.argmax(s)]
+    rect[0] = pts[np.argmin(s)]      # Top-left
+    rect[2] = pts[np.argmax(s)]      # Bottom-right
 
     diff = np.diff(pts, axis=1)
-    rect[1] = pts[np.argmin(diff)]
-    rect[3] = pts[np.argmax(diff)]
+    rect[1] = pts[np.argmin(diff)]   # Top-right
+    rect[3] = pts[np.argmax(diff)]   # Bottom-left
 
-    return rect
+    (tl, tr, br, bl) = rect
 
-# Enhance image lightly
-def enhance_image(img_pil):
-    enhancer = ImageEnhance.Contrast(img_pil)
-    img_pil = enhancer.enhance(1.2)
-    enhancer = ImageEnhance.Sharpness(img_pil)
-    img_pil = enhancer.enhance(1.3)
-    return img_pil
-
-# Save images to a PDF
-def images_to_pdf(images):
-    pdf = FPDF()
-    for img in images:
-        temp_path = "/tmp/temp_img.jpg"
-        img.save(temp_path, "JPEG")
-        width, height = img.size
-        orientation = 'P' if height > width else 'L'
-        pdf.add_page(orientation=orientation)
-        pdf.image(temp_path, x=10, y=10, w=pdf.w-20)
-    pdf_output = "/tmp/cleaned_images.pdf"
-    pdf.output(pdf_output)
-    return pdf_output
-
-# Encode file for download
-def download_link(file_path, file_label):
-    with open(file_path, "rb") as f:
-        base64_pdf = base64.b64encode(f.read()).decode('utf-8')
-    href = f'<a href="data:application/octet-stream;base64,{base64_pdf}" download="{file_label}">📄 Download your scanned PDF</a>'
-    return href
-
-# App Layout
-page_bg_img = '''
-<style>
-.stApp {
-background-image: url("https://images.unsplash.com/photo-1581090700227-1c065c769e3a");
-background-size: cover;
-background-position: top left;
-background-repeat: no-repeat;
-background-attachment: fixed;
-}
-.title {
-    color: white;
-    text-align: center;
-    font-size: 40px;
-    font-weight: bold;
-}
-</style>
-'''
+    widthA = np.linalg.norm(br - bl)
+    widthB = np.linalg.norm(tr - tl)
+    maxWidth = max(int(widthA), int(widthB))
 
-st.markdown(page_bg_img, unsafe_allow_html=True)
-st.markdown('<div class="title">📄 Auto ID Card / Document Scanner</div>', unsafe_allow_html=True)
+    heightA = np.linalg.norm(tr - br)
+    heightB = np.linalg.norm(tl - bl)
+    maxHeight = max(int(heightA), int(heightB))
 
-st.write("### Upload your document or ID card photos 📷")
-uploaded_files = st.file_uploader("Upload multiple files", accept_multiple_files=True, type=["png", "jpg", "jpeg"])
+    dst = np.array([
+        [0, 0],
+        [maxWidth-1, 0],
+        [maxWidth-1, maxHeight-1],
+        [0, maxHeight-1]
+    ], dtype="float32")
 
-if uploaded_files:
-    cleaned_images = []
-    for uploaded_file in uploaded_files:
-        file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
-        img = cv2.imdecode(file_bytes, 1)
+    M = cv2.getPerspectiveTransform(rect, dst)
+    warped = cv2.warpPerspective(img, M, (maxWidth, maxHeight))
 
-        cropped = detect_and_crop(img)
-        cropped_rgb = cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB)
-        img_pil = Image.fromarray(cropped_rgb)
+    return Image.fromarray(warped)
 
-        # Optional light enhancement
-        enhanced_img = enhance_image(img_pil)
+def enhance_image(img_pil):
+    """Enhance the image (brightness and contrast)."""
+    enhancer_brightness = ImageEnhance.Brightness(img_pil)
+    img_bright = enhancer_brightness.enhance(1.2)
 
-        cleaned_images.append(enhanced_img)
+    enhancer_contrast = ImageEnhance.Contrast(img_bright)
+    img_contrast = enhancer_contrast.enhance(1.3)
 
-    st.success(f"✅ Processed {len(cleaned_images)} documents successfully!")
+    return img_contrast
 
-    pdf_file = images_to_pdf(cleaned_images)
+def rotate_image(img_pil, angle):
+    """Rotate image."""
+    return img_pil.rotate(angle, expand=True)
 
-    st.markdown(download_link(pdf_file, "scanned_documents.pdf"), unsafe_allow_html=True)
+# -----------------------
+# Streamlit App
+# -----------------------
 
+# Set app background and style
+page_bg_img = f"""
+<style>
+.stApp {{
+background-image: url("https://images.unsplash.com/photo-1581090700227-1c065c585ab3?ixlib=rb-4.0.3&auto=format&fit=crop&w=1470&q=80");
+background-size: cover;
+background-position: center;
+}}
+h1, h2, h3, h4, h5 {{
+color: white;
+}}
+</style>
+"""
+st.markdown(page_bg_img, unsafe_allow_html=True)
+
+# App Title
+st.title("📄 Smart Document Scanner App")
+
+st.write("Upload your **ID cards or documents**. We'll detect, crop, enhance and generate a clean PDF!")
+
+uploaded_files = st.file_uploader(
+    "Upload Images (JPG, JPEG, PNG supported)", 
+    type=["jpg", "jpeg", "png"], 
+    accept_multiple_files=True
+)
+
+if uploaded_files:
+    final_images = []
+
+    for file in uploaded_files:
+        image = Image.open(file)
+        st.subheader(f"Original Uploaded Image: {file.name}")
+        st.image(image, caption="Original Image", use_column_width=True)
+
+        # Auto crop
+        cropped_img = auto_crop_document(image)
+        st.image(cropped_img, caption="Auto Cropped", use_column_width=True)
+
+        # Rotate option
+        rotate_angle = st.selectbox(f"Rotate {file.name}?", [0, 90, 180, 270], key=f"rotate_{file.name}")
+        if rotate_angle != 0:
+            cropped_img = rotate_image(cropped_img, rotate_angle)
+            st.image(cropped_img, caption=f"Rotated by {rotate_angle}°", use_column_width=True)
+
+        # Enhance option
+        if st.button(f"Enhance {file.name}", key=f"enhance_{file.name}"):
+            enhanced_img = enhance_image(cropped_img)
+            st.image(enhanced_img, caption="Enhanced Image (After)", use_column_width=True)
+            final_images.append(enhanced_img)
+        else:
+            final_images.append(cropped_img)
+
+    # Generate PDF
+    if st.button("📥 Generate PDF"):
+        pdf = FPDF()
+        pdf.set_auto_page_break(0)
+
+        for img in final_images:
+            img_temp = io.BytesIO()
+            img.save(img_temp, format='JPEG')
+            img_temp.seek(0)
+
+            cover = Image.open(img_temp)
+            width, height = cover.size
+
+            # Convert pixel in mm (1px = 0.264583 mm)
+            width, height = float(width * 0.264583), float(height * 0.264583)
+
+            pdf.add_page()
+            pdf.image(img_temp, 0, 0, width, height)
+
+        pdf_output = io.BytesIO()
+        pdf.output(pdf_output)
+        pdf_output.seek(0)
+
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = f"scanned_document_{timestamp}.pdf"
+
+        st.download_button(
+            label="📄 Download Your PDF",
+            data=pdf_output,
+            file_name=filename,
+            mime="application/pdf"
+        )