feat: persistent dataset by name

.gitignore

@@ -2,6 +2,16 @@ venv/
 .env
 __pycache__
 .DS_Store
+
+# Results directory (temporary benchmark results)
+results/
+
+# Dataset thumbnails (too large for git, can be regenerated)
+datasets/*/thumbnails/
+
+# Keep the actual dataset files (golden_labels.json)
+!datasets/*/golden_labels.json
+
+# Legacy data directory (can be removed if no longer used)
 data/
-!data/golden_labels.json
-results/
+!data/golden_labels.json

benchmark.py

@@ -1,4 +1,4 @@
-# benchmark.py (Final Fix)
+# benchmark.py (Updated for Named Datasets)
 
 import os
 import json
@@ -9,18 +9,22 @@ from pathlib import Path
 import math
 
 from geo_bot import GeoBot
-from config import DATA_PATHS, MODELS_CONFIG, SUCCESS_THRESHOLD_KM
+from config import get_data_paths, MODELS_CONFIG, SUCCESS_THRESHOLD_KM
 
 
 class MapGuesserBenchmark:
-    def __init__(self, headless: bool = False):
+    def __init__(self, dataset_name: str = "default", headless: bool = False):
+        self.dataset_name = dataset_name
+        self.data_paths = get_data_paths(dataset_name)
         self.headless = headless
         self.golden_labels = self.load_golden_labels()
-        print(f"📊 Loaded {len(self.golden_labels)} golden label samples")
+        print(
+            f"📊 Loaded {len(self.golden_labels)} samples from dataset '{dataset_name}'"
+        )
 
     def load_golden_labels(self) -> List[Dict]:
         try:
-            with open(DATA_PATHS["golden_labels"], "r") as f:
+            with open(self.data_paths["golden_labels"], "r") as f:
                 return json.load(f).get("samples", [])
         except Exception:
             return []
@@ -75,10 +79,11 @@ class MapGuesserBenchmark:
         **kwargs,
     ) -> Dict:
         if not self.golden_labels:
-            raise ValueError("No golden labels available.")
+            raise ValueError(
+                f"No golden labels available in dataset '{self.dataset_name}'."
+            )
 
         models_to_test = models or list(MODELS_CONFIG.keys())
-        # 使用 max_samples 限制测试样本数量
         num_to_test = (
             min(max_samples, len(self.golden_labels))
             if max_samples is not None
@@ -86,7 +91,7 @@ class MapGuesserBenchmark:
         )
         test_samples = self.golden_labels[:num_to_test]
 
-        print(f"🚀 Starting LIVE benchmark:")
+        print(f"🚀 Starting benchmark on dataset '{self.dataset_name}':")
         print(f"   Models: {models_to_test}")
         print(f"   Samples: {len(test_samples)}")
         print(f"   Temperature: {temperature}")
@@ -105,7 +110,9 @@ class MapGuesserBenchmark:
                     temperature=temperature,
                 ) as bot:
                     for i, sample in enumerate(test_samples):
-                        print('########################################################')
+                        print(
+                            "########################################################"
+                        )
                         print(f"📍 Sample {i + 1}/{len(test_samples)}")
                         try:
                             result = self.run_single_test_with_bot(bot, sample)
@@ -154,9 +161,6 @@ class MapGuesserBenchmark:
 
         bot.controller.setup_clean_environment()
 
-        ## TODO add interactive mode to go ahead, turn around and zoom in/out
-        # Mat still need JS to operate but can use selenium to do it or wrap a MCP server
-
         screenshot = bot.take_screenshot()
         if not screenshot:
             return {
@@ -169,14 +173,11 @@ class MapGuesserBenchmark:
         predicted_lat_lon = bot.analyze_image(screenshot)
         inference_time = time.time() - start_time
 
-        # **核心修复**: 从顶级的 "lat" 和 "lng" 键构造真实坐标字典
         true_coords = {"lat": location_data.get("lat"), "lng": location_data.get("lng")}
 
         true_location = location_data["address"]
         print(f"🔍 True location: {true_location}")
-        # print true coords
         print(f"🔍 True coords: {true_coords}")
-        # print predicted coords
         print(f"🔍 Predicted coords: {predicted_lat_lon}")
         distance_km = self.calculate_distance(true_coords, predicted_lat_lon)
 
@@ -193,16 +194,18 @@ class MapGuesserBenchmark:
         }
 
     def save_results(self, results: List[Dict]):
-        # ... (此函数不变) ...
         if not results:
             return
         try:
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-            results_dir = Path(DATA_PATHS["results"])
+            results_dir = Path(self.data_paths["results"])
             results_dir.mkdir(parents=True, exist_ok=True)
             results_file = results_dir / f"benchmark_results_{timestamp}.json"
             output_data = {
-                "metadata": {"timestamp": datetime.now().isoformat()},
+                "metadata": {
+                    "dataset_name": self.dataset_name,
+                    "timestamp": datetime.now().isoformat(),
+                },
                 "results": results,
             }
             with open(results_file, "w") as f:
@@ -212,7 +215,6 @@ class MapGuesserBenchmark:
             print(f"❌ Error saving results: {e}")
 
     def generate_summary(self, results: List[Dict]) -> Dict:
-        # ... (此函数不变) ...
         summary = {}
         by_model = {}
         for r in results:

config.py

@@ -15,8 +15,17 @@ SELECTORS = {
 # Data collection settings
 DATA_COLLECTION_CONFIG = {
     "wait_after_go": 3,
+    "thumbnail_size": (320, 240),
 }
 
+# Benchmark settings
+BENCHMARK_CONFIG = {
+    "data_collection_samples": 50,
+}
+
+# MapCrunch options
+MAPCRUNCH_OPTIONS = {}
+
 # Model configurations
 MODELS_CONFIG = {
     "gpt-4o": {
@@ -37,8 +46,16 @@ MODELS_CONFIG = {
     },
 }
 
-# Data paths
-DATA_PATHS = {
-    "golden_labels": "data/golden_labels.json",
-    "results": "results/",
-}
+
+# Data paths - now supports named datasets
+def get_data_paths(dataset_name: str = "default"):
+    """Get data paths for a specific dataset"""
+    return {
+        "golden_labels": f"datasets/{dataset_name}/golden_labels.json",
+        "thumbnails": f"datasets/{dataset_name}/thumbnails/",
+        "results": f"results/{dataset_name}/",
+    }
+
+
+# Backward compatibility - default paths
+DATA_PATHS = get_data_paths("default")

data_collector.py

@@ -1,4 +1,4 @@
-# data_collector.py (Final Version for High-Quality Data)
+# data_collector.py (Updated for Named Datasets)
 
 import os
 import json
@@ -12,7 +12,7 @@ from io import BytesIO
 
 from mapcrunch_controller import MapCrunchController
 from config import (
-    DATA_PATHS,
+    get_data_paths,
     BENCHMARK_CONFIG,
     DATA_COLLECTION_CONFIG,
     MAPCRUNCH_OPTIONS,
@@ -20,14 +20,21 @@ from config import (
 
 
 class DataCollector:
-    def __init__(self, headless: bool = False, options: Optional[Dict] = None):
+    def __init__(
+        self,
+        dataset_name: str = "default",
+        headless: bool = False,
+        options: Optional[Dict] = None,
+    ):
+        self.dataset_name = dataset_name
+        self.data_paths = get_data_paths(dataset_name)
         self.controller = MapCrunchController(headless=headless)
         self.data = []
         self.options = options or MAPCRUNCH_OPTIONS
         self.setup_directories()
 
     def setup_directories(self):
-        for path in DATA_PATHS.values():
+        for path in self.data_paths.values():
             if path.endswith("/"):
                 Path(path).mkdir(parents=True, exist_ok=True)
             else:
@@ -37,9 +44,9 @@ class DataCollector:
         self, num_samples: Optional[int] = None, **kwargs
     ) -> List[Dict]:
         num_samples = num_samples or BENCHMARK_CONFIG["data_collection_samples"]
-        print(f"🚀 Starting high-quality data collection for {num_samples} samples...")
-
-        # NOTE: setup_collection_options is not implemented in the provided controller, assuming it's handled manually or not needed.
+        print(
+            f"🚀 Collecting {num_samples} samples for dataset '{self.dataset_name}'..."
+        )
 
         successful_samples = 0
         while successful_samples < num_samples:
@@ -63,57 +70,69 @@ class DataCollector:
         self.save_data()
         return self.data
 
-    # 在 data_collector.py 中替换此函数
-
     def collect_single_location(self) -> Optional[Dict]:
-        """Collects a single location and manually constructs the url_slug."""
+        """Collects a single location with simplified data collection."""
         try:
-            # 1. 获取坐标和标识符
+            # Get coordinates
             coords = self.controller.driver.execute_script(
                 "return { lat: window.panorama.getPosition().lat(), lng: window.panorama.getPosition().lng() };"
             )
             if not coords:
                 raise ValueError("Could not get coordinates.")
 
-            identifiers = self.controller.get_live_location_identifiers()
-            if not identifiers or "pov" not in identifiers:
-                raise ValueError("Could not get POV.")
-
-            address = self.controller.get_current_address()
+            # Get POV data directly from panorama
+            pov_data = self.controller.driver.execute_script("""
+                return {
+                    heading: window.panorama.getPov().heading,
+                    pitch: window.panorama.getPov().pitch,
+                    zoom: window.panorama.getZoom(),
+                    panoId: window.panorama.getPano()
+                };
+            """)
+
+            if not pov_data:
+                raise ValueError("Could not get POV data.")
+
+            # Get address (simplified)
+            address = "Unknown"
+            try:
+                address = self.controller.get_current_address() or "Unknown"
+            except:
+                pass  # Address is optional
 
-            # **2. 核心修复：在Python中手动构建url_slug**
             lat = coords.get("lat")
             lng = coords.get("lng")
-            pov = identifiers.get("pov")
-            # MapCrunch的URL slug中，zoom是0-based，而Google POV是1-based
-            zoom_for_slug = round(pov.get("zoom", 1.0)) - 1
 
-            # 使用 roundNum 函数的逻辑来格式化数字
+            # Simplified URL slug construction
             def round_num(n, d):
                 return f"{n:.{d}f}"
 
+            zoom_for_slug = max(0, round(pov_data.get("zoom", 1.0)) - 1)
             url_slug = (
                 f"{round_num(lat, 6)}_"
                 f"{round_num(lng, 6)}_"
-                f"{round_num(pov.get('heading', 0), 2)}_"
-                f"{round_num(pov.get('pitch', 0) * -1, 2)}_"  # Pitch在slug中是负数
+                f"{round_num(pov_data.get('heading', 0), 2)}_"
+                f"{round_num(pov_data.get('pitch', 0) * -1, 2)}_"
                 f"{zoom_for_slug}"
             )
 
-            # 3. 构建数据样本
             sample_id = str(uuid.uuid4())
             location_data = {
                 "id": sample_id,
                 "timestamp": datetime.now().isoformat(),
                 "lat": lat,
                 "lng": lng,
-                "address": address or "Unknown",
-                "pano_id": identifiers.get("panoId"),
-                "pov": pov,
-                "url_slug": url_slug,  # <-- 现在这里永远有正确的值
+                "address": address,
+                "pano_id": pov_data.get("panoId"),
+                "pov": {
+                    "heading": pov_data.get("heading", 0),
+                    "pitch": pov_data.get("pitch", 0),
+                    "zoom": pov_data.get("zoom", 1.0),
+                },
+                "url_slug": url_slug,
             }
 
-            # 4. 保存缩略图
+            # Try to save thumbnail (optional)
             thumbnail_path = self.save_thumbnail(sample_id)
             if thumbnail_path:
                 location_data["thumbnail_path"] = thumbnail_path
@@ -124,38 +143,50 @@ class DataCollector:
             print(f"❌ Error in collect_single_location: {e}")
             return None
 
-    # ... (save_thumbnail, save_data 等其他函数保持不变) ...
     def save_thumbnail(self, sample_id: str) -> Optional[str]:
         try:
             screenshot_bytes = self.controller.take_street_view_screenshot()
             if not screenshot_bytes:
+                print(
+                    f"⚠️  Could not take screenshot for {sample_id} (this is OK in headless mode)"
+                )
                 return None
+
             image = Image.open(BytesIO(screenshot_bytes))
             thumbnail_size = DATA_COLLECTION_CONFIG.get("thumbnail_size", (320, 240))
             image.thumbnail(thumbnail_size, Image.Resampling.LANCZOS)
             thumbnail_filename = f"{sample_id}.jpg"
-            thumbnail_path = os.path.join(DATA_PATHS["thumbnails"], thumbnail_filename)
+            thumbnail_path = os.path.join(
+                self.data_paths["thumbnails"], thumbnail_filename
+            )
+
             if image.mode in ("RGBA", "LA"):
                 rgb_image = Image.new("RGB", image.size, (255, 255, 255))
                 rgb_image.paste(image, mask=image.split()[-1])
                 image = rgb_image
+
             image.save(thumbnail_path, "JPEG", quality=85)
+            print(f"✅ Saved thumbnail for {sample_id}")
             return thumbnail_filename
-        except Exception:
+        except Exception as e:
+            print(f"⚠️  Could not save thumbnail for {sample_id}: {e}")
             return None
 
     def save_data(self):
         try:
             output_data = {
                 "metadata": {
+                    "dataset_name": self.dataset_name,
                     "collection_date": datetime.now().isoformat(),
                     "collection_options": self.options,
                 },
                 "samples": self.data,
             }
-            with open(DATA_PATHS["golden_labels"], "w") as f:
+            with open(self.data_paths["golden_labels"], "w") as f:
                 json.dump(output_data, f, indent=2)
-            print(f"\n💾 High-quality data saved to {DATA_PATHS['golden_labels']}")
+            print(
+                f"\n💾 Dataset '{self.dataset_name}' saved to {self.data_paths['golden_labels']}"
+            )
         except Exception as e:
             print(f"❌ Error saving data: {e}")
 

datasets/asia/golden_labels.json

@@ -0,0 +1,159 @@
+{
+  "metadata": {
+    "dataset_name": "asia",
+    "collection_date": "2025-06-11T21:13:45.005091",
+    "collection_options": {}
+  },
+  "samples": [
+    {
+      "id": "fdbb9997-c07c-4d4d-9095-82f162f0c27a",
+      "timestamp": "2025-06-11T21:13:15.310368",
+      "lat": 42.1322878067665,
+      "lng": 26.787410093767097,
+      "address": "Unknown",
+      "pano_id": "gsRLllGBndoh4EMBklXL9Q",
+      "pov": {
+        "heading": 240.93000000000006,
+        "pitch": 5,
+        "zoom": 1.0000051533649421
+      },
+      "url_slug": "42.132288_26.787410_240.93_-5.00_0",
+      "thumbnail_path": "fdbb9997-c07c-4d4d-9095-82f162f0c27a.jpg"
+    },
+    {
+      "id": "05ac7262-54b6-4b81-b9be-85d0830d7ed1",
+      "timestamp": "2025-06-11T21:13:18.720025",
+      "lat": 50.815127864126566,
+      "lng": 3.3295800788973042,
+      "address": "Unknown",
+      "pano_id": "lw4NuJ2I82JRsk5y8N7gGA",
+      "pov": {
+        "heading": -262.06999999999994,
+        "pitch": 5,
+        "zoom": 1.0000070740241276
+      },
+      "url_slug": "50.815128_3.329580_-262.07_-5.00_0",
+      "thumbnail_path": "05ac7262-54b6-4b81-b9be-85d0830d7ed1.jpg"
+    },
+    {
+      "id": "9fff9f32-e6f0-4f5f-a476-d019ec8b5bf2",
+      "timestamp": "2025-06-11T21:13:21.994062",
+      "lat": 45.43514041007389,
+      "lng": 21.364097624705536,
+      "address": "Unknown",
+      "pano_id": "XLiyeDvQ9SoaSaDBxn3GDA",
+      "pov": {
+        "heading": -168.06999999999994,
+        "pitch": 5,
+        "zoom": 1.000006919588194
+      },
+      "url_slug": "45.435140_21.364098_-168.07_-5.00_0",
+      "thumbnail_path": "9fff9f32-e6f0-4f5f-a476-d019ec8b5bf2.jpg"
+    },
+    {
+      "id": "2da4a9b9-e597-46e9-8c9d-8701f0c63462",
+      "timestamp": "2025-06-11T21:13:25.252476",
+      "lat": 42.30902518065906,
+      "lng": 77.8748629197877,
+      "address": "Unknown",
+      "pano_id": "t0HJFo38t3rh1U6W2OZ_VA",
+      "pov": {
+        "heading": 5.930000000000064,
+        "pitch": 5,
+        "zoom": 1.0000092331912114
+      },
+      "url_slug": "42.309025_77.874863_5.93_-5.00_0",
+      "thumbnail_path": "2da4a9b9-e597-46e9-8c9d-8701f0c63462.jpg"
+    },
+    {
+      "id": "b93858db-454e-4cc7-8f61-a3fe09cb0cab",
+      "timestamp": "2025-06-11T21:13:28.569143",
+      "lat": 14.647613688319248,
+      "lng": -16.980851505792,
+      "address": "Unknown",
+      "pano_id": "GE8DtAXvn2qZuSALopw8xA",
+      "pov": {
+        "heading": 270.93000000000006,
+        "pitch": 5,
+        "zoom": 1.0000015226193344
+      },
+      "url_slug": "14.647614_-16.980852_270.93_-5.00_0",
+      "thumbnail_path": "b93858db-454e-4cc7-8f61-a3fe09cb0cab.jpg"
+    },
+    {
+      "id": "1f0e20f1-3687-4939-be23-7c7b490cc707",
+      "timestamp": "2025-06-11T21:13:31.763851",
+      "lat": 11.208463091095442,
+      "lng": 105.72569729813453,
+      "address": "Unknown",
+      "pano_id": "2W3x5T-dMOrMJO57YtGq2Q",
+      "pov": {
+        "heading": -316.06999999999994,
+        "pitch": 5,
+        "zoom": 1.000004517085056
+      },
+      "url_slug": "11.208463_105.725697_-316.07_-5.00_0",
+      "thumbnail_path": "1f0e20f1-3687-4939-be23-7c7b490cc707.jpg"
+    },
+    {
+      "id": "7a2049cc-959c-4948-8574-0ffe9950e86e",
+      "timestamp": "2025-06-11T21:13:35.017287",
+      "lat": 40.668879231679576,
+      "lng": -8.21452809466328,
+      "address": "Unknown",
+      "pano_id": "bYHqbKkNgUUffaYf6fcKBQ",
+      "pov": {
+        "heading": -86.06999999999994,
+        "pitch": 5,
+        "zoom": 1.0000035631232127
+      },
+      "url_slug": "40.668879_-8.214528_-86.07_-5.00_0",
+      "thumbnail_path": "7a2049cc-959c-4948-8574-0ffe9950e86e.jpg"
+    },
+    {
+      "id": "fec89220-b130-49b3-8c19-f7e2e9551acf",
+      "timestamp": "2025-06-11T21:13:38.277525",
+      "lat": 5.459666786657994,
+      "lng": -2.1135681235966626,
+      "address": "Unknown",
+      "pano_id": "AaoI1zHJ4Pf18j94UXqksA",
+      "pov": {
+        "heading": 122.93000000000006,
+        "pitch": 5,
+        "zoom": 1.0000018586016313
+      },
+      "url_slug": "5.459667_-2.113568_122.93_-5.00_0",
+      "thumbnail_path": "fec89220-b130-49b3-8c19-f7e2e9551acf.jpg"
+    },
+    {
+      "id": "0fd8d569-c98a-4af3-b08c-173121c76043",
+      "timestamp": "2025-06-11T21:13:41.520505",
+      "lat": 1.3921814170475024,
+      "lng": 103.98320353936504,
+      "address": "Unknown",
+      "pano_id": "CAoSF0NJSE0wb2dLRUlDQWdJQ2t3T1h4bWdF",
+      "pov": {
+        "heading": 278.93000000000006,
+        "pitch": 5,
+        "zoom": 1.0000078101439185
+      },
+      "url_slug": "1.392181_103.983204_278.93_-5.00_0",
+      "thumbnail_path": "0fd8d569-c98a-4af3-b08c-173121c76043.jpg"
+    },
+    {
+      "id": "2426fad0-aeda-426a-96ad-f16724c360ce",
+      "timestamp": "2025-06-11T21:13:44.746699",
+      "lat": 51.89795854217673,
+      "lng": 4.96480321921333,
+      "address": "Unknown",
+      "pano_id": "bxBt_sZjG7ocUqPSmD1X0Q",
+      "pov": {
+        "heading": -175.06999999999994,
+        "pitch": 5,
+        "zoom": 1.0000093258133977
+      },
+      "url_slug": "51.897959_4.964803_-175.07_-5.00_0",
+      "thumbnail_path": "2426fad0-aeda-426a-96ad-f16724c360ce.jpg"
+    }
+  ]
+}

list_datasets.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""
+Utility script to list available datasets
+"""
+
+import json
+import os
+from pathlib import Path
+from config import get_data_paths
+
+
+def list_datasets():
+    """List all available datasets"""
+    datasets_dir = Path("datasets")
+    if not datasets_dir.exists():
+        print("No datasets directory found.")
+        return []
+
+    datasets = []
+    for dataset_dir in datasets_dir.iterdir():
+        if dataset_dir.is_dir():
+            dataset_name = dataset_dir.name
+            data_paths = get_data_paths(dataset_name)
+            golden_labels_path = data_paths["golden_labels"]
+
+            if os.path.exists(golden_labels_path):
+                try:
+                    with open(golden_labels_path, "r") as f:
+                        data = json.load(f)
+                        samples = data.get("samples", [])
+                        metadata = data.get("metadata", {})
+
+                    datasets.append(
+                        {
+                            "name": dataset_name,
+                            "samples": len(samples),
+                            "created": metadata.get("collection_date", "Unknown"),
+                            "path": golden_labels_path,
+                        }
+                    )
+                except Exception as e:
+                    print(f"❌ Error reading dataset '{dataset_name}': {e}")
+
+    return datasets
+
+
+def main():
+    print("📊 Available Datasets:")
+    print("=" * 50)
+
+    datasets = list_datasets()
+
+    if not datasets:
+        print("No datasets found.")
+        print("\nTo create a new dataset, run:")
+        print("python main.py --mode collect --dataset <name> --samples <count>")
+        return
+
+    for dataset in sorted(datasets, key=lambda x: x["name"]):
+        print(f"Dataset: {dataset['name']}")
+        print(f"  Samples: {dataset['samples']}")
+        print(f"  Created: {dataset['created']}")
+        print(f"  Path: {dataset['path']}")
+        print()
+
+    print("To use a dataset, run:")
+    print("python main.py --mode benchmark --dataset <name>")
+    print("python main.py --mode agent --dataset <name>")
+
+
+if __name__ == "__main__":
+    main()

main.py

@@ -9,49 +9,62 @@ from langchain_google_genai import ChatGoogleGenerativeAI
 
 from geo_bot import GeoBot
 from benchmark import MapGuesserBenchmark
-from config import MODELS_CONFIG, DATA_PATHS, SUCCESS_THRESHOLD_KM
-
-
-def agent_mode(model_name: str, steps: int, headless: bool, samples: int, temperature: float = 0.0):
+from data_collector import DataCollector
+from config import MODELS_CONFIG, get_data_paths, SUCCESS_THRESHOLD_KM
+
+
+def agent_mode(
+    model_name: str,
+    steps: int,
+    headless: bool,
+    samples: int,
+    dataset_name: str = "default",
+    temperature: float = 0.0,
+):
     """
     Runs the AI Agent in a benchmark loop over multiple samples,
     using multi-step exploration for each.
     """
     print(
-        f"Starting Agent Mode (as a benchmark): model={model_name}, steps={steps}, samples={samples}, temperature={temperature}"
+        f"Starting Agent Mode: model={model_name}, steps={steps}, samples={samples}, dataset={dataset_name}, temperature={temperature}"
     )
 
+    data_paths = get_data_paths(dataset_name)
     try:
-        with open(DATA_PATHS["golden_labels"], "r", encoding="utf-8") as f:
+        with open(data_paths["golden_labels"], "r", encoding="utf-8") as f:
             golden_labels = json.load(f).get("samples", [])
     except FileNotFoundError:
-        print(f"Error: Golden labels file not found at {DATA_PATHS['golden_labels']}.")
+        print(
+            f"Error: Dataset '{dataset_name}' not found at {data_paths['golden_labels']}."
+        )
         return
 
     if not golden_labels:
-        print("Error: No samples found in golden_labels.json.")
+        print(f"Error: No samples found in dataset '{dataset_name}'.")
         return
 
     num_to_test = min(samples, len(golden_labels))
     test_samples = golden_labels[:num_to_test]
-    print(f"Will run on {len(test_samples)} samples.")
+    print(f"Will run on {len(test_samples)} samples from dataset '{dataset_name}'.")
 
     config = MODELS_CONFIG.get(model_name)
     model_class = globals()[config["class"]]
     model_instance_name = config["model_name"]
 
-    benchmark_helper = MapGuesserBenchmark(headless=True)
+    benchmark_helper = MapGuesserBenchmark(dataset_name=dataset_name, headless=True)
     all_results = []
 
     with GeoBot(
-        model=model_class, model_name=model_instance_name, headless=headless, temperature=temperature
+        model=model_class,
+        model_name=model_instance_name,
+        headless=headless,
+        temperature=temperature,
     ) as bot:
         for i, sample in enumerate(test_samples):
             print(
                 f"\n--- Running Sample {i + 1}/{len(test_samples)} (ID: {sample.get('id')}) ---"
             )
 
-            # **FIXED**: Correct sequence: Load Data -> Clean Environment -> Run Loop
             if not bot.controller.load_location_from_data(sample):
                 print(
                     f"   ❌ Failed to load location for sample {sample.get('id')}. Skipping."
@@ -98,36 +111,61 @@ def agent_mode(model_name: str, steps: int, headless: bool, samples: int, temper
 
     summary = benchmark_helper.generate_summary(all_results)
     if summary:
-        print("\n\n--- Agent Benchmark Complete! Summary ---")
+        print(
+            f"\n\n--- Agent Benchmark Complete for dataset '{dataset_name}'! Summary ---"
+        )
         for model, stats in summary.items():
             print(f"Model: {model}")
             print(f"  Success Rate: {stats['success_rate'] * 100:.1f}%")
             print(f"  Avg Distance: {stats['average_distance_km']:.1f} km")
 
-    print("\nAgent Mode finished.")
+    print("Agent Mode finished.")
 
 
-def benchmark_mode(models: list, samples: int, headless: bool, temperature: float = 0.0):
+def benchmark_mode(
+    models: list,
+    samples: int,
+    headless: bool,
+    dataset_name: str = "default",
+    temperature: float = 0.0,
+):
     """Runs the benchmark on pre-collected data."""
-    print(f"Starting Benchmark Mode: models={models}, samples={samples}, temperature={temperature}")
-    benchmark = MapGuesserBenchmark(headless=headless)
-    summary = benchmark.run_benchmark(models=models, max_samples=samples, temperature=temperature)
+    print(
+        f"Starting Benchmark Mode: models={models}, samples={samples}, dataset={dataset_name}, temperature={temperature}"
+    )
+    benchmark = MapGuesserBenchmark(dataset_name=dataset_name, headless=headless)
+    summary = benchmark.run_benchmark(
+        models=models, max_samples=samples, temperature=temperature
+    )
     if summary:
-        print("\n--- Benchmark Complete! Summary ---")
+        print(f"\n--- Benchmark Complete for dataset '{dataset_name}'! Summary ---")
         for model, stats in summary.items():
             print(f"Model: {model}")
             print(f"  Success Rate: {stats['success_rate'] * 100:.1f}%")
             print(f"  Avg Distance: {stats['average_distance_km']:.1f} km")
 
 
+def collect_mode(dataset_name: str, samples: int, headless: bool):
+    """Collects data for a new dataset."""
+    print(f"Starting Data Collection: dataset={dataset_name}, samples={samples}")
+    with DataCollector(dataset_name=dataset_name, headless=headless) as collector:
+        collector.collect_samples(num_samples=samples)
+    print(f"Data collection complete for dataset '{dataset_name}'.")
+
+
 def main():
     parser = argparse.ArgumentParser(description="MapCrunch AI Agent & Benchmark")
     parser.add_argument(
         "--mode",
-        choices=["agent", "benchmark"],
+        choices=["agent", "benchmark", "collect"],
         default="agent",
         help="Operation mode.",
     )
+    parser.add_argument(
+        "--dataset",
+        default="default",
+        help="Dataset name to use or create.",
+    )
     parser.add_argument(
         "--model",
         choices=list(MODELS_CONFIG.keys()),
@@ -161,12 +199,19 @@ def main():
 
     args = parser.parse_args()
 
-    if args.mode == "agent":
+    if args.mode == "collect":
+        collect_mode(
+            dataset_name=args.dataset,
+            samples=args.samples,
+            headless=args.headless,
+        )
+    elif args.mode == "agent":
         agent_mode(
             model_name=args.model,
             steps=args.steps,
             headless=args.headless,
             samples=args.samples,
+            dataset_name=args.dataset,
             temperature=args.temperature,
         )
     elif args.mode == "benchmark":
@@ -174,6 +219,7 @@ def main():
             models=args.models or [args.model],
             samples=args.samples,
             headless=args.headless,
+            dataset_name=args.dataset,
             temperature=args.temperature,
         )