Spaces:
Sleeping
Sleeping
#!/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() |