Update app.py

app.py

@@ -3,9 +3,10 @@ import os
 import glob
 import time
 import streamlit as st
+import fitz  # PyMuPDF
+import requests
 from PIL import Image
-import torch
-from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, AutoTokenizer, AutoModel, TrOCRProcessor, VisionEncoderDecoderModel
+from transformers import AutoTokenizer, AutoModel
 from diffusers import StableDiffusionPipeline
 import cv2
 import numpy as np
@@ -31,21 +32,18 @@ st.set_page_config(
     page_icon="🤖",
     layout="wide",
     initial_sidebar_state="expanded",
-    menu_items={'About': "AI Vision Titans: OCR, Image Gen, Line Drawings on CPU! 🌌"}
+    menu_items={'About': "AI Vision Titans: PDF Snapshots, OCR, Image Gen, Line Drawings on CPU! 🌌"}
 )
 
 # Initialize st.session_state
-if 'captured_images' not in st.session_state:
-    st.session_state['captured_images'] = []
+if 'captured_files' not in st.session_state:
+    st.session_state['captured_files'] = []
 if 'processing' not in st.session_state:
     st.session_state['processing'] = {}
 
 # Utility Functions
 def generate_filename(sequence, ext="png"):
-    from datetime import datetime
-    import pytz
-    central = pytz.timezone('US/Central')
-    timestamp = datetime.now(central).strftime("%d%m%Y%H%M%S%p")
+    timestamp = time.strftime("%d%m%Y%H%M%S")
     return f"{sequence}{timestamp}.{ext}"
 
 def get_gallery_files(file_types):
@@ -61,20 +59,27 @@ def update_gallery():
                     st.image(Image.open(file), caption=file, use_container_width=True)
                 elif file.endswith(".txt"):
                     with open(file, "r") as f:
-                        st.text(f.read()[:50] + "..." if len(f.read()) > 50 else f.read(), help=file)
+                        content = f.read()
+                        st.text(content[:50] + "..." if len(content) > 50 else content, help=file)
 
-# Model Loaders (Smaller, CPU-focused)
-def load_ocr_qwen2vl():
-    model_id = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"
-    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-    model = Qwen2VLForConditionalGeneration.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
-    return processor, model
+def download_pdf(url, output_path):
+    try:
+        response = requests.get(url, stream=True, timeout=10)
+        if response.status_code == 200:
+            with open(output_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            return True
+    except requests.RequestException as e:
+        logger.error(f"Failed to download {url}: {e}")
+    return False
 
-def load_ocr_trocr():
-    model_id = "microsoft/trocr-small-handwritten"  # ~250 MB
-    processor = TrOCRProcessor.from_pretrained(model_id)
-    model = VisionEncoderDecoderModel.from_pretrained(model_id, torch_dtype=torch.float32).to("cpu").eval()
-    return processor, model
+# Model Loaders (CPU-focused)
+def load_ocr_got():
+    model_id = "ucaslcl/GOT-OCR2_0"
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    model = AutoModel.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
+    return tokenizer, model
 
 def load_image_gen():
     model_id = "OFA-Sys/small-stable-diffusion-v0"  # ~300 MB
@@ -82,37 +87,60 @@ def load_image_gen():
     return pipeline
 
 def load_line_drawer():
-    # Simplified OpenCV-based edge detection (CPU-friendly substitute for Torch Space UNet)
-    def edge_detection(image):
+    def edge_detection(image, style="fine"):
         img_np = np.array(image.convert("RGB"))
         gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
-        edges = cv2.Canny(gray, 100, 200)
+        if style == "fine":
+            edges = cv2.Canny(gray, 50, 150)  # Finer lines
+        else:  # "bold"
+            edges = cv2.Canny(gray, 100, 200)  # Bolder lines
         return Image.fromarray(edges)
     return edge_detection
 
 # Async Processing Functions
-async def process_ocr(image, prompt, model_name, output_file):
+async def process_pdf_snapshot(pdf_path, mode="thumbnail"):
     start_time = time.time()
     status = st.empty()
-    status.text(f"Processing {model_name} OCR... (0s)")
-    if model_name == "Qwen2-VL-OCR-2B":
-        processor, model = load_ocr_qwen2vl()
-        # Corrected input format: apply chat template
-        messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt}]}]
-        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[text], images=[image], return_tensors="pt", padding=True).to("cpu")
-        outputs = model.generate(**inputs, max_new_tokens=1024)
-        result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-    else:  # TrOCR
-        processor, model = load_ocr_trocr()
-        pixel_values = processor(images=image, return_tensors="pt").pixel_values.to("cpu")
-        outputs = model.generate(pixel_values)
-        result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+    status.text(f"Processing PDF Snapshot ({mode})... (0s)")
+    doc = fitz.open(pdf_path)
+    output_files = []
+    
+    if mode == "thumbnail":
+        page = doc[0]
+        pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))  # 50% scale
+        output_file = generate_filename("thumbnail", "png")
+        pix.save(output_file)
+        output_files.append(output_file)
+    elif mode == "twopage":
+        for i in range(min(2, len(doc))):
+            page = doc[i]
+            pix = page.get_pixmap(matrix=fitz.Matrix(1.0, 1.0))  # Full scale
+            output_file = generate_filename(f"twopage_{i}", "png")
+            pix.save(output_file)
+            output_files.append(output_file)
+    
+    doc.close()
     elapsed = int(time.time() - start_time)
-    status.text(f"{model_name} OCR completed in {elapsed}s!")
+    status.text(f"PDF Snapshot ({mode}) completed in {elapsed}s!")
+    for file in output_files:
+        if file not in st.session_state['captured_files']:
+            st.session_state['captured_files'].append(file)
+    update_gallery()
+    return output_files
+
+async def process_ocr(image, output_file):
+    start_time = time.time()
+    status = st.empty()
+    status.text("Processing GOT-OCR2_0... (0s)")
+    tokenizer, model = load_ocr_got()
+    result = model.chat(tokenizer, image, ocr_type='ocr')
+    elapsed = int(time.time() - start_time)
+    status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
     async with aiofiles.open(output_file, "w") as f:
         await f.write(result)
-    st.session_state['captured_images'].append(output_file)
+    if output_file not in st.session_state['captured_files']:
+        st.session_state['captured_files'].append(output_file)
+    update_gallery()
     return result
 
 async def process_image_gen(prompt, output_file):
@@ -120,30 +148,34 @@ async def process_image_gen(prompt, output_file):
     status = st.empty()
     status.text("Processing Image Gen... (0s)")
     pipeline = load_image_gen()
-    gen_image = pipeline(prompt, num_inference_steps=20).images[0]  # Reduced steps for speed
+    gen_image = pipeline(prompt, num_inference_steps=20).images[0]
     elapsed = int(time.time() - start_time)
     status.text(f"Image Gen completed in {elapsed}s!")
     gen_image.save(output_file)
-    st.session_state['captured_images'].append(output_file)
+    if output_file not in st.session_state['captured_files']:
+        st.session_state['captured_files'].append(output_file)
+    update_gallery()
     return gen_image
 
-async def process_line_drawing(image, output_file):
+async def process_line_drawing(image, style, output_file):
     start_time = time.time()
     status = st.empty()
-    status.text("Processing Line Drawing... (0s)")
+    status.text(f"Processing Line Drawing ({style})... (0s)")
     edge_fn = load_line_drawer()
-    line_drawing = edge_fn(image)
+    line_drawing = edge_fn(image, style=style)
     elapsed = int(time.time() - start_time)
-    status.text(f"Line Drawing completed in {elapsed}s!")
+    status.text(f"Line Drawing ({style}) completed in {elapsed}s!")
     line_drawing.save(output_file)
-    st.session_state['captured_images'].append(output_file)
+    if output_file not in st.session_state['captured_files']:
+        st.session_state['captured_files'].append(output_file)
+    update_gallery()
     return line_drawing
 
 # Main App
-st.title("AI Vision Titans 🚀 (OCR, Gen, Drawings!)")
+st.title("AI Vision Titans 🚀")
 
 # Sidebar Gallery
-st.sidebar.header("Captured Images 🎨")
+st.sidebar.header("Captured Files 📜")
 gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 4)
 update_gallery()
 
@@ -154,7 +186,7 @@ with log_container:
         st.write(f"{record.asctime} - {record.levelname} - {record.message}")
 
 # Tabs
-tab1, tab2, tab3, tab4 = st.tabs(["Camera Snap 📷", "Test OCR 🔍", "Test Image Gen 🎨", "Test Line Drawings ✏️"])
+tab1, tab2, tab3, tab4, tab5 = st.tabs(["Camera Snap 📷", "Download PDFs 📥", "Test OCR 🔍", "Test Image Gen 🎨", "Test Line Drawings ✏️"])
 
 with tab1:
     st.header("Camera Snap 📷")
@@ -164,23 +196,23 @@ with tab1:
         cam0_img = st.camera_input("Take a picture - Cam 0", key="cam0")
         if cam0_img:
             filename = generate_filename(0)
-            if filename not in st.session_state['captured_images']:
+            if filename not in st.session_state['captured_files']:
                 with open(filename, "wb") as f:
                     f.write(cam0_img.getvalue())
                 st.image(Image.open(filename), caption=filename, use_container_width=True)
                 logger.info(f"Saved snapshot from Camera 0: {filename}")
-                st.session_state['captured_images'].append(filename)
+                st.session_state['captured_files'].append(filename)
                 update_gallery()
     with cols[1]:
         cam1_img = st.camera_input("Take a picture - Cam 1", key="cam1")
         if cam1_img:
             filename = generate_filename(1)
-            if filename not in st.session_state['captured_images']:
+            if filename not in st.session_state['captured_files']:
                 with open(filename, "wb") as f:
                     f.write(cam1_img.getvalue())
                 st.image(Image.open(filename), caption=filename, use_container_width=True)
                 logger.info(f"Saved snapshot from Camera 1: {filename}")
-                st.session_state['captured_images'].append(filename)
+                st.session_state['captured_files'].append(filename)
                 update_gallery()
 
     st.subheader("Burst Capture")
@@ -194,42 +226,57 @@ with tab1:
                 img = st.camera_input(f"Frame {i}", key=f"burst_{i}_{time.time()}")
                 if img:
                     filename = generate_filename(f"burst_{i}")
-                    if filename not in st.session_state['captured_images']:
+                    if filename not in st.session_state['captured_files']:
                         with open(filename, "wb") as f:
                             f.write(img.getvalue())
                         st.session_state['burst_frames'].append(filename)
                         logger.info(f"Saved burst frame {i}: {filename}")
                         st.image(Image.open(filename), caption=filename, use_container_width=True)
-                    time.sleep(0.5)  # Small delay for visibility
-        st.session_state['captured_images'].extend([f for f in st.session_state['burst_frames'] if f not in st.session_state['captured_images']])
+                    time.sleep(0.5)
+        st.session_state['captured_files'].extend([f for f in st.session_state['burst_frames'] if f not in st.session_state['captured_files']])
         update_gallery()
         placeholder.success(f"Captured {len(st.session_state['burst_frames'])} frames!")
 
 with tab2:
+    st.header("Download PDFs 📥")
+    url_input = st.text_area("Enter PDF URLs (one per line)", height=100)
+    mode = st.selectbox("Snapshot Mode", ["Thumbnail", "Two-Page View"], key="download_mode")
+    if st.button("Download & Snapshot 📸"):
+        urls = url_input.strip().split("\n")
+        for url in urls:
+            if url:
+                pdf_path = generate_filename("downloaded", "pdf")
+                if download_pdf(url, pdf_path):
+                    logger.info(f"Downloaded PDF from {url} to {pdf_path}")
+                    snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode.lower().replace(" ", "")))
+                    for snapshot in snapshots:
+                        st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
+                else:
+                    st.error(f"Failed to download {url}")
+
+with tab3:
     st.header("Test OCR 🔍")
-    captured_images = get_gallery_files(["png"])
-    if captured_images:
-        selected_image = st.selectbox("Select Image", captured_images, key="ocr_select")
-        image = Image.open(selected_image)
+    captured_files = get_gallery_files(["png"])
+    if captured_files:
+        selected_file = st.selectbox("Select Image", captured_files, key="ocr_select")
+        image = Image.open(selected_file)
         st.image(image, caption="Input Image", use_container_width=True)
-        ocr_model = st.selectbox("Select OCR Model", ["Qwen2-VL-OCR-2B", "TrOCR-Small"], key="ocr_model_select")
-        prompt = st.text_area("Prompt", "Extract text from the image", key="ocr_prompt")
         if st.button("Run OCR 🚀", key="ocr_run"):
             output_file = generate_filename("ocr_output", "txt")
             st.session_state['processing']['ocr'] = True
-            result = asyncio.run(process_ocr(image, prompt, ocr_model, output_file))
+            result = asyncio.run(process_ocr(image, output_file))
             st.text_area("OCR Result", result, height=200, key="ocr_result")
             st.success(f"OCR output saved to {output_file}")
             st.session_state['processing']['ocr'] = False
     else:
-        st.warning("No images captured yet. Use Camera Snap first!")
+        st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")
 
-with tab3:
+with tab4:
     st.header("Test Image Gen 🎨")
-    captured_images = get_gallery_files(["png"])
-    if captured_images:
-        selected_image = st.selectbox("Select Image", captured_images, key="gen_select")
-        image = Image.open(selected_image)
+    captured_files = get_gallery_files(["png"])
+    if captured_files:
+        selected_file = st.selectbox("Select Image", captured_files, key="gen_select")
+        image = Image.open(selected_file)
         st.image(image, caption="Reference Image", use_container_width=True)
         prompt = st.text_area("Prompt", "Generate a similar superhero image", key="gen_prompt")
         if st.button("Run Image Gen 🚀", key="gen_run"):
@@ -240,24 +287,25 @@ with tab3:
             st.success(f"Image saved to {output_file}")
             st.session_state['processing']['gen'] = False
     else:
-        st.warning("No images captured yet. Use Camera Snap first!")
+        st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")
 
-with tab4:
+with tab5:
     st.header("Test Line Drawings ✏️")
-    captured_images = get_gallery_files(["png"])
-    if captured_images:
-        selected_image = st.selectbox("Select Image", captured_images, key="line_select")
-        image = Image.open(selected_image)
+    captured_files = get_gallery_files(["png"])
+    if captured_files:
+        selected_file = st.selectbox("Select Image", captured_files, key="line_select")
+        image = Image.open(selected_file)
         st.image(image, caption="Input Image", use_container_width=True)
+        style = st.selectbox("Line Style", ["Fine", "Bold"], key="line_style")
         if st.button("Run Line Drawing 🚀", key="line_run"):
-            output_file = generate_filename("line_output", "png")
+            output_file = generate_filename(f"line_{style.lower()}", "png")
             st.session_state['processing']['line'] = True
-            result = asyncio.run(process_line_drawing(image, output_file))
-            st.image(result, caption="Line Drawing", use_container_width=True)
+            result = asyncio.run(process_line_drawing(image, style.lower(), output_file))
+            st.image(result, caption=f"{style} Line Drawing", use_container_width=True)
             st.success(f"Line drawing saved to {output_file}")
             st.session_state['processing']['line'] = False
     else:
-        st.warning("No images captured yet. Use Camera Snap first!")
+        st.warning("No images captured yet. Use Camera Snap or Download PDFs first!")
 
 # Initial Gallery Update
 update_gallery()