backup5.app.py

#!/usr/bin/env python3
import os
import glob
import time
import streamlit as st
from PIL import Image
import torch
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, AutoTokenizer, AutoModel, TrOCRProcessor, VisionEncoderDecoderModel
from diffusers import StableDiffusionPipeline
import cv2
import numpy as np
import logging
import asyncio
import aiofiles
from io import BytesIO

# Logging setup
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
log_records = []

class LogCaptureHandler(logging.Handler):
    def emit(self, record):
        log_records.append(record)

logger.addHandler(LogCaptureHandler())

# Page Configuration
st.set_page_config(
    page_title="AI Vision Titans 🚀",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={'About': "AI Vision Titans: 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 '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")
    return f"{sequence}{timestamp}.{ext}"

def get_gallery_files(file_types):
    return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])

def update_gallery():
    media_files = get_gallery_files(["png", "txt"])
    if media_files:
        cols = st.sidebar.columns(2)
        for idx, file in enumerate(media_files[:gallery_size * 2]):
            with cols[idx % 2]:
                if file.endswith(".png"):
                    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)

# 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 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

def load_image_gen():
    model_id = "OFA-Sys/small-stable-diffusion-v0"  # ~300 MB
    pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32).to("cpu")
    return pipeline

def load_line_drawer():
    # Simplified OpenCV-based edge detection (CPU-friendly substitute for Torch Space UNet)
    def edge_detection(image):
        img_np = np.array(image.convert("RGB"))
        gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
        edges = cv2.Canny(gray, 100, 200)
        return Image.fromarray(edges)
    return edge_detection

# Async Processing Functions
async def process_ocr(image, prompt, model_name, output_file):
    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]
    elapsed = int(time.time() - start_time)
    status.text(f"{model_name} OCR 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)
    return result

async def process_image_gen(prompt, output_file):
    start_time = time.time()
    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
    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)
    return gen_image

async def process_line_drawing(image, output_file):
    start_time = time.time()
    status = st.empty()
    status.text("Processing Line Drawing... (0s)")
    edge_fn = load_line_drawer()
    line_drawing = edge_fn(image)
    elapsed = int(time.time() - start_time)
    status.text(f"Line Drawing completed in {elapsed}s!")
    line_drawing.save(output_file)
    st.session_state['captured_images'].append(output_file)
    return line_drawing

# Main App
st.title("AI Vision Titans 🚀 (OCR, Gen, Drawings!)")

# Sidebar Gallery
st.sidebar.header("Captured Images 🎨")
gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 4)
update_gallery()

st.sidebar.subheader("Action Logs 📜")
log_container = st.sidebar.empty()
with log_container:
    for record in log_records:
        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 ✏️"])

with tab1:
    st.header("Camera Snap 📷")
    st.subheader("Single Capture")
    cols = st.columns(2)
    with cols[0]:
        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']:
                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)
                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']:
                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)
                update_gallery()

    st.subheader("Burst Capture")
    slice_count = st.number_input("Number of Frames", min_value=1, max_value=20, value=10, key="burst_count")
    if st.button("Start Burst Capture 📸"):
        st.session_state['burst_frames'] = []
        placeholder = st.empty()
        for i in range(slice_count):
            with placeholder.container():
                st.write(f"Capturing frame {i+1}/{slice_count}...")
                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']:
                        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']])
        update_gallery()
        placeholder.success(f"Captured {len(st.session_state['burst_frames'])} frames!")

with tab2:
    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)
        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))
            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!")

with tab3:
    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)
        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"):
            output_file = generate_filename("gen_output", "png")
            st.session_state['processing']['gen'] = True
            result = asyncio.run(process_image_gen(prompt, output_file))
            st.image(result, caption="Generated Image", use_container_width=True)
            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!")

with tab4:
    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)
        st.image(image, caption="Input Image", use_container_width=True)
        if st.button("Run Line Drawing 🚀", key="line_run"):
            output_file = generate_filename("line_output", "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)
            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!")

# Initial Gallery Update
update_gallery()