Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import json | |
| import os | |
| import time | |
| from io import BytesIO | |
| from PIL import Image | |
| from pathlib import Path | |
| from geo_bot import GeoBot, AGENT_PROMPT_TEMPLATE | |
| from benchmark import MapGuesserBenchmark | |
| from config import MODELS_CONFIG, get_data_paths, SUCCESS_THRESHOLD_KM | |
| from langchain_openai import ChatOpenAI | |
| from langchain_anthropic import ChatAnthropic | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from hf_chat import HuggingFaceChat | |
| # Simple API key setup | |
| if "OPENAI_API_KEY" in st.secrets: | |
| os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"] | |
| if "ANTHROPIC_API_KEY" in st.secrets: | |
| os.environ["ANTHROPIC_API_KEY"] = st.secrets["ANTHROPIC_API_KEY"] | |
| if "GOOGLE_API_KEY" in st.secrets: | |
| os.environ["GOOGLE_API_KEY"] = st.secrets["GOOGLE_API_KEY"] | |
| if "HF_TOKEN" in st.secrets: | |
| os.environ["HF_TOKEN"] = st.secrets["HF_TOKEN"] | |
| def get_available_datasets(): | |
| datasets_dir = Path("datasets") | |
| if not datasets_dir.exists(): | |
| return ["default"] | |
| datasets = [] | |
| for dataset_dir in datasets_dir.iterdir(): | |
| if dataset_dir.is_dir(): | |
| data_paths = get_data_paths(dataset_dir.name) | |
| if os.path.exists(data_paths["golden_labels"]): | |
| datasets.append(dataset_dir.name) | |
| return datasets if datasets else ["default"] | |
| def get_model_class(class_name): | |
| if class_name == "ChatOpenAI": | |
| return ChatOpenAI | |
| elif class_name == "ChatAnthropic": | |
| return ChatAnthropic | |
| elif class_name == "ChatGoogleGenerativeAI": | |
| return ChatGoogleGenerativeAI | |
| elif class_name == "HuggingFaceChat": | |
| return HuggingFaceChat | |
| else: | |
| raise ValueError(f"Unknown model class: {class_name}") | |
| # UI Setup | |
| st.set_page_config(page_title="MapCrunch AI Agent", layout="wide") | |
| st.title("🗺️ MapCrunch AI Agent") | |
| # Sidebar | |
| with st.sidebar: | |
| st.header("⚙️ Configuration") | |
| dataset_choice = st.selectbox("Dataset", get_available_datasets()) | |
| model_choice = st.selectbox("Model", list(MODELS_CONFIG.keys())) | |
| steps_per_sample = st.slider("Max Steps", 3, 20, 10) | |
| # Load dataset | |
| data_paths = get_data_paths(dataset_choice) | |
| with open(data_paths["golden_labels"], "r") as f: | |
| golden_labels = json.load(f).get("samples", []) | |
| st.info(f"Dataset has {len(golden_labels)} samples") | |
| num_samples = st.slider( | |
| "Samples to Test", 1, len(golden_labels), min(3, len(golden_labels)) | |
| ) | |
| start_button = st.button("🚀 Start", type="primary") | |
| # Main Logic | |
| if start_button: | |
| test_samples = golden_labels[:num_samples] | |
| config = MODELS_CONFIG[model_choice] | |
| model_class = get_model_class(config["class"]) | |
| benchmark_helper = MapGuesserBenchmark(dataset_name=dataset_choice) | |
| all_results = [] | |
| progress_bar = st.progress(0) | |
| with GeoBot( | |
| model=model_class, model_name=config["model_name"], headless=True | |
| ) as bot: | |
| for i, sample in enumerate(test_samples): | |
| st.divider() | |
| st.header(f"Sample {i + 1}/{num_samples}") | |
| bot.controller.load_location_from_data(sample) | |
| bot.controller.setup_clean_environment() | |
| col1, col2 = st.columns([2, 3]) | |
| with col1: | |
| image_placeholder = st.empty() | |
| with col2: | |
| reasoning_placeholder = st.empty() | |
| action_placeholder = st.empty() | |
| history = [] | |
| final_guess = None | |
| for step in range(steps_per_sample): | |
| step_num = step + 1 | |
| reasoning_placeholder.info(f"🤔 Step {step_num}/{steps_per_sample}") | |
| bot.controller.label_arrows_on_screen() | |
| screenshot_bytes = bot.controller.take_street_view_screenshot() | |
| image_placeholder.image(screenshot_bytes, caption=f"Step {step_num}") | |
| current_step = { | |
| "image_b64": bot.pil_to_base64( | |
| Image.open(BytesIO(screenshot_bytes)) | |
| ), | |
| "action": "N/A", | |
| } | |
| history.append(current_step) | |
| available_actions = bot.controller.get_available_actions() | |
| history_text = "\n".join( | |
| [f"Step {j + 1}: {h['action']}" for j, h in enumerate(history[:-1])] | |
| ) | |
| if not history_text: | |
| history_text = "First step." | |
| prompt = AGENT_PROMPT_TEMPLATE.format( | |
| remaining_steps=steps_per_sample - step, | |
| history_text=history_text, | |
| available_actions=json.dumps(available_actions), | |
| ) | |
| message = bot._create_message_with_history( | |
| prompt, [h["image_b64"] for h in history] | |
| ) | |
| response = bot.model.invoke(message) | |
| decision = bot._parse_agent_response(response) | |
| if not decision: | |
| decision = { | |
| "action_details": {"action": "PAN_RIGHT"}, | |
| "reasoning": "Fallback", | |
| } | |
| action = decision.get("action_details", {}).get("action") | |
| history[-1]["action"] = action | |
| reasoning_placeholder.success("✅ Decision Made") | |
| action_placeholder.success(f"🎯 Action: `{action}`") | |
| with action_placeholder: | |
| with st.expander("Reasoning"): | |
| st.write(decision.get("reasoning", "N/A")) | |
| if step_num == steps_per_sample and action != "GUESS": | |
| action = "GUESS" | |
| if action == "GUESS": | |
| lat = decision.get("action_details", {}).get("lat") | |
| lon = decision.get("action_details", {}).get("lon") | |
| if lat is not None and lon is not None: | |
| final_guess = (lat, lon) | |
| break | |
| elif action == "MOVE_FORWARD": | |
| bot.controller.move("forward") | |
| elif action == "MOVE_BACKWARD": | |
| bot.controller.move("backward") | |
| elif action == "PAN_LEFT": | |
| bot.controller.pan_view("left") | |
| elif action == "PAN_RIGHT": | |
| bot.controller.pan_view("right") | |
| time.sleep(1) | |
| # Results | |
| true_coords = {"lat": sample.get("lat"), "lng": sample.get("lng")} | |
| distance_km = None | |
| is_success = False | |
| if final_guess: | |
| distance_km = benchmark_helper.calculate_distance( | |
| true_coords, final_guess | |
| ) | |
| if distance_km is not None: | |
| is_success = distance_km <= SUCCESS_THRESHOLD_KM | |
| st.subheader("🎯 Result") | |
| col1, col2, col3 = st.columns(3) | |
| col1.metric("Guess", f"{final_guess[0]:.3f}, {final_guess[1]:.3f}") | |
| col2.metric( | |
| "Truth", f"{true_coords['lat']:.3f}, {true_coords['lng']:.3f}" | |
| ) | |
| col3.metric( | |
| "Distance", | |
| f"{distance_km:.1f} km", | |
| delta="Success" if is_success else "Failed", | |
| ) | |
| all_results.append( | |
| { | |
| "sample_id": sample.get("id"), | |
| "model": model_choice, | |
| "true_coordinates": true_coords, | |
| "predicted_coordinates": final_guess, | |
| "distance_km": distance_km, | |
| "success": is_success, | |
| } | |
| ) | |
| progress_bar.progress((i + 1) / num_samples) | |
| # Summary | |
| st.divider() | |
| st.header("🏁 Summary") | |
| summary = benchmark_helper.generate_summary(all_results) | |
| if summary and model_choice in summary: | |
| stats = summary[model_choice] | |
| col1, col2 = st.columns(2) | |
| col1.metric("Success Rate", f"{stats.get('success_rate', 0) * 100:.1f}%") | |
| col2.metric("Avg Distance", f"{stats.get('average_distance_km', 0):.1f} km") | |
| st.dataframe(all_results) | |