#!/usr/bin/env python3 import os import streamlit as st from PIL import Image import torch from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, AutoTokenizer, AutoModel from diffusers import StableDiffusionPipeline import cv2 import numpy as np import logging 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'] = [] # 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"]) if media_files: cols = st.sidebar.columns(2) for idx, file in enumerate(media_files[:gallery_size * 2]): with cols[idx % 2]: st.image(Image.open(file), caption=file, use_container_width=True) # Model Loaders (Simplified, 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_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" # Small, CPU-friendly pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32).to("cpu") return pipeline def load_line_drawer(): # Simplified from your Torch Space (assuming a UNet-like model for edge detection) # Placeholder: Using OpenCV edge detection as a minimal CPU example 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 # 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 📷") slice_count = st.number_input("Image Slice Count", min_value=1, max_value=20, value=10) cols = st.columns(2) with cols[0]: st.subheader("Camera 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() if st.button(f"Capture {slice_count} Frames - Cam 0 📸"): st.session_state['cam0_frames'] = [] for i in range(slice_count): img = st.camera_input(f"Frame {i} - Cam 0", key=f"cam0_frame_{i}_{time.time()}") if img: filename = generate_filename(f"0_{i}") if filename not in st.session_state['captured_images']: with open(filename, "wb") as f: f.write(img.getvalue()) st.session_state['cam0_frames'].append(filename) logger.info(f"Saved frame {i} from Camera 0: {filename}") time.sleep(1.0 / slice_count) st.session_state['captured_images'].extend([f for f in st.session_state['cam0_frames'] if f not in st.session_state['captured_images']]) update_gallery() for frame in st.session_state['cam0_frames']: st.image(Image.open(frame), caption=frame, use_container_width=True) with cols[1]: st.subheader("Camera 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() if st.button(f"Capture {slice_count} Frames - Cam 1 📸"): st.session_state['cam1_frames'] = [] for i in range(slice_count): img = st.camera_input(f"Frame {i} - Cam 1", key=f"cam1_frame_{i}_{time.time()}") if img: filename = generate_filename(f"1_{i}") if filename not in st.session_state['captured_images']: with open(filename, "wb") as f: f.write(img.getvalue()) st.session_state['cam1_frames'].append(filename) logger.info(f"Saved frame {i} from Camera 1: {filename}") time.sleep(1.0 / slice_count) st.session_state['captured_images'].extend([f for f in st.session_state['cam1_frames'] if f not in st.session_state['captured_images']]) update_gallery() for frame in st.session_state['cam1_frames']: st.image(Image.open(frame), caption=frame, use_container_width=True) with tab2: st.header("Test OCR 🔍") captured_images = get_gallery_files(["png"]) if captured_images: selected_image = st.selectbox("Select Image", captured_images) 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", "GOT-OCR2_0"]) prompt = st.text_area("Prompt", "Extract text from the image") if st.button("Run OCR 🚀"): if ocr_model == "Qwen2-VL-OCR-2B": processor, model = load_ocr_qwen2vl() inputs = processor(text=[prompt], images=[image], return_tensors="pt").to("cpu") outputs = model.generate(**inputs, max_new_tokens=1024) text = processor.decode(outputs[0], skip_special_tokens=True) else: # GOT-OCR2_0 tokenizer, model = load_ocr_got() with open(selected_image, "rb") as f: img_bytes = f.read() img = Image.open(BytesIO(img_bytes)) text = model.chat(tokenizer, img, ocr_type='ocr') st.text_area("OCR Result", text, height=200) with tab3: st.header("Test Image Gen 🎨") captured_images = get_gallery_files(["png"]) if captured_images: selected_image = st.selectbox("Select Image", captured_images) 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") if st.button("Run Image Gen 🚀"): pipeline = load_image_gen() gen_image = pipeline(prompt, num_inference_steps=50).images[0] st.image(gen_image, caption="Generated Image", use_container_width=True) with tab4: st.header("Test Line Drawings ✏️") captured_images = get_gallery_files(["png"]) if captured_images: selected_image = st.selectbox("Select Image", captured_images) image = Image.open(selected_image) st.image(image, caption="Input Image", use_container_width=True) if st.button("Run Line Drawing 🚀"): edge_fn = load_line_drawer() line_drawing = edge_fn(image) st.image(line_drawing, caption="Line Drawing", use_container_width=True) # Initial Gallery Update update_gallery()