import streamlit as st import time import utils from PIL import Image import numpy as np import uuid # Set page config st.set_page_config(page_title="Annotation Assistant", layout="wide", page_icon="✨") # --- Premium Custom CSS --- st.markdown(""" """, unsafe_allow_html=True) # --- State Management --- if "model_loaded" not in st.session_state: st.session_state.model_loaded = False if "sessions" not in st.session_state: # Structure: { session_id: { name, history, detections, image, metrics, timestamp } } st.session_state.sessions = {} if "active_session_id" not in st.session_state: st.session_state.active_session_id = None # Helper 1: Create a new session def create_session(name="New Chat"): session_id = str(uuid.uuid4()) st.session_state.sessions[session_id] = { "name": name, "history": [], "detections": [], "image": None, "metrics": {}, "created_at": time.time() } st.session_state.active_session_id = session_id return session_id # Helper 2: Get active session data def get_active_session(): if not st.session_state.active_session_id: create_session() return st.session_state.sessions[st.session_state.active_session_id] # Ensure at least one session exists if not st.session_state.sessions: create_session() current_session = get_active_session() # --- Sidebar (Session Manager) --- with st.sidebar: st.markdown("### 🗂ī¸ Sessions") if st.button("➕ New Chat", use_container_width=True, type="primary"): create_session() st.rerun() st.markdown("---") # Sort sessions by recency sorted_sessions = sorted( st.session_state.sessions.items(), key=lambda x: x[1]['created_at'], reverse=True ) for s_id, s_data in sorted_sessions: # Hide empty "New Chat" sessions from the list unless active if s_data['image'] is None: continue is_active = (s_id == st.session_state.active_session_id) display_name = s_data['name'] icon = "📂" if is_active else "📝" label = f"{icon} {display_name}" if st.button(label, key=f"sess_{s_id}", use_container_width=True, type="secondary" if not is_active else "primary"): st.session_state.active_session_id = s_id st.rerun() # --- Model Loading --- if not st.session_state.model_loaded: with st.spinner("Initializing AI Core..."): processor, model = utils.load_model() if processor and model: st.session_state.model_loaded = True st.session_state.processor = processor st.session_state.model = model st.rerun() else: st.error("Model Engine Failure.") st.stop() # --- Main Workspace --- # Header col_logo, col_space = st.columns([6, 1]) with col_logo: if current_session['name'] == "New Chat": st.markdown("# Annotation Assistant") else: st.markdown(f"# {current_session['name']}") # Logic if current_session['image'] is None: # --- Upload State --- st.markdown( "

Upload an image to start this session

", unsafe_allow_html=True ) uploaded_file = st.file_uploader( "Upload Image", type=["jpg", "png", "jpeg"], key=f"uploader_{st.session_state.active_session_id}", label_visibility="collapsed" ) if uploaded_file: image = Image.open(uploaded_file).convert("RGB") current_session['image'] = image current_session['name'] = uploaded_file.name st.rerun() else: # --- Analysis State --- # Image Controls img_width = st.slider("Adjust View Size", 300, 1500, 700, 50, help="Drag to resize the image view") st.markdown("
", unsafe_allow_html=True) # 1. Main visual (Hero) display_image = current_session['image'].copy() if current_session['detections']: display_image = utils.draw_boxes(display_image, current_session['detections']) st.image(display_image, width=img_width) # 2. Results Actions & Metrics if current_session['detections']: # Metrics Row if current_session['metrics']: m = current_session['metrics'] st.markdown(f"""
Inference {m.get('inference_time', 0)}s | Total {m.get('total_time', 0)}s | Tokens {m.get('token_count', 0)}
""", unsafe_allow_html=True) # Download Row c1, c2, c3 = st.columns([1, 1, 3]) # Bias to left with c1: # UPDATED: Pass usage metadata for Strict COCO compatibility coco_json = utils.convert_to_coco( current_session['detections'], image_size=current_session['image'].size, filename=current_session['name'] ) st.download_button("Download JSON", coco_json, "annotations.json", "application/json", use_container_width=True) with c2: zip_buffer = utils.create_crops_zip(current_session['image'], current_session['detections']) st.download_button("Download ZIP", zip_buffer, "crops.zip", "application/zip", use_container_width=True) # 3. Reasoning Stream (Below) st.markdown("
", unsafe_allow_html=True) st.markdown("### AI Insights") with st.container(): st.markdown("
", unsafe_allow_html=True) for det in current_session['detections'][::-1]: label = det.get('label', 'Object') reasoning = det.get('reasoning', None) if not reasoning: reasoning = "Object detected based on visual features." st.markdown(f"""
{label}
{reasoning}
""", unsafe_allow_html=True) st.markdown("
", unsafe_allow_html=True) else: # Image loaded but no detections st.markdown( "
" "Waiting for instructions... Use the chat bar below." "
", unsafe_allow_html=True ) # --- Floating Chat Bar --- st.markdown("
", unsafe_allow_html=True) prompt = st.chat_input("Describe objects to detect...") if prompt: if current_session['image'] is None: st.error("Please upload an image first.") else: # Warning for HF Spaces Free Tier if "cpu" in str(st.session_state.get('device', 'cpu')): st.info("ℹī¸ Running on CPU (Free Tier). Complex scenes may take 30-60s to analyze.") with st.status("Analyzing Scene...", expanded=True) as status: detections, updated_history, raw_text, metrics = utils.get_bounding_boxes( current_session['image'], prompt, current_session['history'], st.session_state.processor, st.session_state.model ) if detections: current_session['detections'] = utils.smart_merge_detections(current_session['detections'], detections) current_session['history'] = updated_history current_session['metrics'] = metrics status.update(label="Complete", state="complete", expanded=False) st.rerun() else: status.update(label="No matches found.", state="error", expanded=False) st.toast(f"No match found.", icon="⚠ī¸")