Initial VPR demo implementation

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,14 +1,44 @@
----
-title: Simple Vpr Demo
-emoji: 🐠
-colorFrom: green
-colorTo: gray
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
-license: mit
-short_description: 'This space, is a simple demo for a vpr system. '
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Visual Place Recognition Demo
+
+This is a Visual Place Recognition (VPR) demo using EigenPlaces model. Upload a query image to find similar places in our database of 400+ images from various cities.
+
+## How it works
+
+1. Upload a query image
+2. The model extracts visual features from your image
+3. It compares these features with pre-computed features from 400+ database images
+4. Returns the most similar matches with similarity scores and location information
+
+## Dataset
+
+- **Database**: 400+ images from various cities
+- **Cities**: Melbourne, Boston, and others
+- **Metadata**: Each image includes place ID, city, and GPS coordinates
+
+## Model
+
+- **Architecture**: EigenPlaces with ResNet50 backbone
+- **Descriptor Dimension**: 2048
+- **Similarity Metric**: Cosine similarity
+
+## Usage
+
+1. Upload a query image using the interface
+2. Adjust the number of matches you want to see (1-10)
+3. Click "Find Matches" to get results
+4. View the matched images and their metadata
+
+## Technical Details
+
+The demo uses:
+- EigenPlaces model for visual feature extraction
+- Pre-computed descriptors for fast similarity search
+- Cosine similarity for matching
+- Gradio for the web interface
+
+## Files
+
+- `app.py`: Main Gradio application
+- `model.py`: Model loading and descriptor computation
+- `dataset.py`: Dataset handling and ground truth lookup
+- `data/`: Contains database images, query images, and ground truth JSON

app.py

@@ -0,0 +1,154 @@
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+import json
+from pathlib import Path
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from model import load_model
+from dataset import VPRDataset
+
+# Global variables
+model = None
+dataset = None
+db_descriptors = None
+db_filenames = None
+
+def load_everything():
+    """Load model, dataset, and pre-compute database descriptors."""
+    global model, dataset, db_descriptors, db_filenames
+    
+    print("Loading model...")
+    model = load_model("eigenplaces")
+    
+    print("Loading dataset...")
+    dataset = VPRDataset('data')
+    
+    print("Pre-computing database descriptors...")
+    # Create database-only dataset
+    db_dataset = VPRDataset('data', include_queries=False, include_database=True)
+    
+    # Create DataLoader for efficient batch processing
+    batch_size = 1
+    dataloader = DataLoader(db_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
+    
+    # Compute descriptors for database images using DataLoader
+    db_descriptors = []
+    db_filenames = []
+    
+    model.eval()
+    with torch.no_grad():
+        for batch_images, batch_filenames, batch_is_query in tqdm(dataloader, desc="Computing database descriptors"):
+            # Move batch to same device as model
+            device = next(model.parameters()).device
+            batch_images = batch_images.to(device)
+            
+            # Compute descriptors for this batch
+            batch_descriptors = model(batch_images)
+            
+            # Store results
+            db_descriptors.append(batch_descriptors.cpu())
+            db_filenames.extend(batch_filenames)
+    
+    # Concatenate all descriptors
+    db_descriptors = torch.cat(db_descriptors, dim=0)
+    print(f"Pre-computed descriptors for {len(db_filenames)} database images")
+
+def find_matches(query_image, top_k=5):
+    """Find top-k matches for a query image."""
+    if model is None or db_descriptors is None:
+        return "Model not loaded yet. Please wait..."
+    
+    # Extract query descriptor
+    with torch.no_grad():
+        query_batch = query_image.unsqueeze(0)
+        query_descriptor = model(query_batch).cpu()
+    
+    # Compute similarities (cosine similarity)
+    query_norm = query_descriptor / torch.norm(query_descriptor)
+    db_norm = db_descriptors / torch.norm(db_descriptors, dim=1, keepdim=True)
+    similarities = torch.mm(query_norm, db_norm.T).squeeze()
+    
+    # Get top-k matches
+    top_similarities, top_indices = torch.topk(similarities, top_k)
+    
+    # Prepare results
+    results = []
+    for i, (sim, idx) in enumerate(zip(top_similarities, top_indices)):
+        filename = db_filenames[idx]
+        img_path = Path('data') / 'database' / filename
+        
+        # Get metadata
+        item_info = dataset.get_item_by_filename(filename)
+        
+        results.append({
+            'image': str(img_path),
+            'similarity': float(sim),
+            'place_id': item_info['place_id'],
+            'city': item_info['city'],
+            'coordinates': f"{item_info['lat']:.4f}, {item_info['lon']:.4f}"
+        })
+    
+    return results
+
+def demo_interface(query_image, top_k):
+    """Gradio interface function."""
+    if query_image is None:
+        return "Please upload a query image"
+    
+    # Convert PIL to tensor format expected by model
+    import torchvision.transforms as T
+    transform = T.Compose([
+        T.Resize((480, 640)),
+        T.ToTensor(),
+        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    
+    query_tensor = transform(query_image)
+    matches = find_matches(query_tensor, int(top_k))
+    
+    if isinstance(matches, str):
+        return matches
+    
+    # Format results for display
+    result_text = "Top Matches:\n\n"
+    result_images = []
+    
+    for i, match in enumerate(matches):
+        result_text += f"{i+1}. Similarity: {match['similarity']:.4f}\n"
+        result_text += f"   Place ID: {match['place_id']}\n"
+        result_text += f"   City: {match['city']}\n"
+        result_text += f"   Coordinates: {match['coordinates']}\n\n"
+        
+        result_images.append(match['image'])
+    
+    return result_text, result_images
+
+# Load everything on startup
+load_everything()
+
+# Create Gradio interface
+with gr.Blocks(title="Visual Place Recognition Demo") as demo:
+    gr.Markdown("# Visual Place Recognition Demo")
+    gr.Markdown("Upload a query image to find similar places in our database!")
+    
+    with gr.Row():
+        with gr.Column():
+            query_input = gr.Image(type="pil", label="Query Image")
+            top_k_slider = gr.Slider(1, 10, value=5, step=1, label="Number of matches")
+            find_button = gr.Button("Find Matches")
+        
+        with gr.Column():
+            result_text = gr.Textbox(label="Results", lines=10)
+            result_gallery = gr.Gallery(label="Matched Images", show_label=True)
+    
+    find_button.click(
+        fn=demo_interface,
+        inputs=[query_input, top_k_slider],
+        outputs=[result_text, result_gallery]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

data/ground_truth.json

@@ -0,0 +1,182 @@
+{
+  "database": [
+    {
+      "filename": "place00004796_db0000.jpg",
+      "place_id": 4796,
+      "city": "London",
+      "lat": 51.50916209875723,
+      "lon": -0.1489464938656002
+    },
+    {
+      "filename": "place00004796_db0001.jpg",
+      "place_id": 4796,
+      "city": "London",
+      "lat": 51.5091715210003,
+      "lon": -0.148933838176506
+    },
+    {
+      "filename": "place00004796_db0002.jpg",
+      "place_id": 4796,
+      "city": "London",
+      "lat": 51.50917252788888,
+      "lon": -0.1489267203772275
+    },
+    {
+      "filename": "place00004796_db0003.jpg",
+      "place_id": 4796,
+      "city": "London",
+      "lat": 51.50916624802473,
+      "lon": -0.1489330344067381
+    },
+    {
+      "filename": "place00203981_db0004.jpg",
+      "place_id": 203981,
+      "city": "Melbourne",
+      "lat": -37.81509362211957,
+      "lon": 144.992838452311
+    },
+    {
+      "filename": "place00203981_db0005.jpg",
+      "place_id": 203981,
+      "city": "Melbourne",
+      "lat": -37.81506767547648,
+      "lon": 144.9928412577368
+    },
+    {
+      "filename": "place00203981_db0006.jpg",
+      "place_id": 203981,
+      "city": "Melbourne",
+      "lat": -37.81509508467622,
+      "lon": 144.9928784011848
+    },
+    {
+      "filename": "place00203981_db0007.jpg",
+      "place_id": 203981,
+      "city": "Melbourne",
+      "lat": -37.81507084333941,
+      "lon": 144.9928344766533
+    },
+    {
+      "filename": "place00008797_db0008.jpg",
+      "place_id": 8797,
+      "city": "London",
+      "lat": 51.52845407460239,
+      "lon": -0.1750029953860952
+    },
+    {
+      "filename": "place00008797_db0009.jpg",
+      "place_id": 8797,
+      "city": "London",
+      "lat": 51.52846288309365,
+      "lon": -0.1750064077374475
+    },
+    {
+      "filename": "place00008797_db0010.jpg",
+      "place_id": 8797,
+      "city": "London",
+      "lat": 51.52845549538673,
+      "lon": -0.1750071427715894
+    },
+    {
+      "filename": "place00008797_db0011.jpg",
+      "place_id": 8797,
+      "city": "London",
+      "lat": 51.52842807679626,
+      "lon": -0.1750499767058685
+    },
+    {
+      "filename": "place00201236_db0012.jpg",
+      "place_id": 201236,
+      "city": "Melbourne",
+      "lat": -37.84278510997986,
+      "lon": 144.9907996152359
+    },
+    {
+      "filename": "place00201236_db0013.jpg",
+      "place_id": 201236,
+      "city": "Melbourne",
+      "lat": -37.84281467344245,
+      "lon": 144.9908206918488
+    },
+    {
+      "filename": "place00201236_db0014.jpg",
+      "place_id": 201236,
+      "city": "Melbourne",
+      "lat": -37.8428116471166,
+      "lon": 144.9907948010103
+    },
+    {
+      "filename": "place00201236_db0015.jpg",
+      "place_id": 201236,
+      "city": "Melbourne",
+      "lat": -37.84279737925578,
+      "lon": 144.9907950980963
+    },
+    {
+      "filename": "place00205527_db0016.jpg",
+      "place_id": 205527,
+      "city": "Melbourne",
+      "lat": -37.79846599314756,
+      "lon": 144.9649501082595
+    },
+    {
+      "filename": "place00205527_db0017.jpg",
+      "place_id": 205527,
+      "city": "Melbourne",
+      "lat": -37.79846121519402,
+      "lon": 144.9649625622069
+    },
+    {
+      "filename": "place00205527_db0018.jpg",
+      "place_id": 205527,
+      "city": "Melbourne",
+      "lat": -37.79846527742957,
+      "lon": 144.9650067670199
+    },
+    {
+      "filename": "place00205527_db0019.jpg",
+      "place_id": 205527,
+      "city": "Melbourne",
+      "lat": -37.79846435981845,
+      "lon": 144.964976424101
+    }
+  ],
+  "query": [
+    {
+      "filename": "place00004796_q0000.jpg",
+      "place_id": 4796,
+      "city": "London",
+      "lat": 51.5091691701721,
+      "lon": -0.1489371501232689
+    },
+    {
+      "filename": "place00203981_q0001.jpg",
+      "place_id": 203981,
+      "city": "Melbourne",
+      "lat": -37.81508165748332,
+      "lon": 144.9928762644381
+    },
+    {
+      "filename": "place00008797_q0002.jpg",
+      "place_id": 8797,
+      "city": "London",
+      "lat": 51.52844848114152,
+      "lon": -0.1750700055099293
+    },
+    {
+      "filename": "place00201236_q0003.jpg",
+      "place_id": 201236,
+      "city": "Melbourne",
+      "lat": -37.84279084700875,
+      "lon": 144.9907951180407
+    },
+    {
+      "filename": "place00205527_q0004.jpg",
+      "place_id": 205527,
+      "city": "Melbourne",
+      "lat": -37.79845501432714,
+      "lon": 144.9649401788737
+    }
+  ],
+  "place_mapping": {}
+}

dataset.py

@@ -0,0 +1,313 @@
+"""
+Simple PyTorch Dataset for VPR (Visual Place Recognition)
+Combines database and query images with ground truth lookup.
+"""
+
+import json
+from pathlib import Path
+from PIL import Image
+import torch
+from torch.utils.data import Dataset
+import torchvision.transforms as T
+from typing import List, Dict, Tuple
+
+
+class VPRDataset(Dataset):
+    """
+    Simple VPR Dataset that loads both database and query images.
+    
+    Usage:
+        dataset = VPRDataset('data')
+        
+        # Get an image
+        img, filename, is_query = dataset[0]
+        
+        # Get ground truth matches for a query
+        matches = dataset.gt('place00000123_q0000.jpg')
+    """
+    
+    def __init__(
+        self, 
+        data_dir='data',
+        transform=None,
+        include_queries=True,
+        include_database=True
+    ):
+        """
+        Args:
+            data_dir: Path to data folder containing database/, query/, and ground_truth.json
+            transform: Optional torchvision transforms to apply to images
+            include_queries: Whether to include query images in the dataset
+            include_database: Whether to include database images in the dataset
+        """
+        self.data_dir = Path(data_dir)
+        self.include_queries = include_queries
+        self.include_database = include_database
+        
+        # Default transform if none provided
+        if transform is None:
+            self.transform = T.Compose([
+                T.Resize((480, 640)),
+                T.ToTensor(),
+                T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+            ])
+        else:
+            self.transform = transform
+        
+        # Load ground truth
+        gt_path = self.data_dir / 'ground_truth.json'
+        with open(gt_path, 'r') as f:
+            self.ground_truth = json.load(f)
+        
+        # Build the dataset items list
+        self.items = []
+        
+        if include_database:
+            for item in self.ground_truth['database']:
+                self.items.append({
+                    'filename': item['filename'],
+                    'path': self.data_dir / 'database' / item['filename'],
+                    'place_id': item['place_id'],
+                    'is_query': False,
+                    'city': item['city'],
+                    'lat': item['lat'],
+                    'lon': item['lon']
+                })
+        
+        if include_queries:
+            for item in self.ground_truth['query']:
+                self.items.append({
+                    'filename': item['filename'],
+                    'path': self.data_dir / 'query' / item['filename'],
+                    'place_id': item['place_id'],
+                    'is_query': True,
+                    'city': item['city'],
+                    'lat': item['lat'],
+                    'lon': item['lon']
+                })
+        
+        # Build lookup tables for fast ground truth queries
+        self._build_lookup_tables()
+    
+    def _build_lookup_tables(self):
+        """Build internal lookup tables for efficient ground truth queries."""
+        # Map filename -> full item info
+        self.filename_to_item = {item['filename']: item for item in self.items}
+        
+        # Map place_id -> list of database filenames
+        self.place_to_db_files = {}
+        for item in self.ground_truth['database']:
+            place_id = item['place_id']
+            if place_id not in self.place_to_db_files:
+                self.place_to_db_files[place_id] = []
+            self.place_to_db_files[place_id].append(item['filename'])
+        
+        # Map query filename -> its place_id for fast lookup
+        self.query_to_place = {
+            item['filename']: item['place_id'] 
+            for item in self.ground_truth['query']
+        }
+    
+    def __len__(self):
+        """Return total number of images in dataset."""
+        return len(self.items)
+    
+    def __getitem__(self, idx) -> Tuple[torch.Tensor, str, bool]:
+        """
+        Get an image from the dataset.
+        
+        Returns:
+            tuple: (image_tensor, filename, is_query)
+                - image_tensor: Transformed image as torch.Tensor
+                - filename: String filename (e.g., 'place00000123_db0001.jpg')
+                - is_query: Boolean indicating if this is a query image
+        """
+        item = self.items[idx]
+        
+        # Load image
+        img = Image.open(item['path']).convert('RGB')
+        
+        # Apply transforms
+        if self.transform:
+            img = self.transform(img)
+        
+        return img, item['filename'], item['is_query']
+    
+    def gt(self, query_filename: str) -> List[str]:
+        """
+        Get ground truth database matches for a query image.
+        
+        Args:
+            query_filename: Filename of the query image (e.g., 'place00000123_q0000.jpg')
+        
+        Returns:
+            List of database image filenames that match this query (same place_id)
+        
+        Example:
+            >>> dataset = VPRDataset('data')
+            >>> matches = dataset.gt('place00000123_q0000.jpg')
+            >>> print(matches)
+            ['place00000123_db0000.jpg', 'place00000123_db0001.jpg', 'place00000123_db0002.jpg']
+        """
+        if query_filename not in self.query_to_place:
+            raise ValueError(f"Query filename '{query_filename}' not found in dataset")
+        
+        place_id = self.query_to_place[query_filename]
+        return self.place_to_db_files.get(place_id, [])
+    
+    def get_query_filenames(self) -> List[str]:
+        """Get list of all query image filenames."""
+        return list(self.query_to_place.keys())
+    
+    def get_database_filenames(self) -> List[str]:
+        """Get list of all database image filenames."""
+        all_db_files = []
+        for files in self.place_to_db_files.values():
+            all_db_files.extend(files)
+        return all_db_files
+    
+    def get_item_by_filename(self, filename: str) -> Dict:
+        """
+        Get full item information by filename.
+        
+        Args:
+            filename: Image filename
+        
+        Returns:
+            Dictionary with keys: filename, path, place_id, is_query, city, lat, lon
+        """
+        if filename not in self.filename_to_item:
+            raise ValueError(f"Filename '{filename}' not found in dataset")
+        return self.filename_to_item[filename]
+    
+    @staticmethod
+    def get_place_id_from_filename(filename: str) -> int:
+        """
+        Extract place_id from filename.
+        
+        Args:
+            filename: Image filename (e.g., 'place00000123_db0001.jpg')
+        
+        Returns:
+            Integer place_id (e.g., 123)
+        """
+        return int(filename.split('_')[0].replace('place', ''))
+
+
+# ============================================================================
+# EXAMPLE USAGE
+# ============================================================================
+
+if __name__ == "__main__":
+    from torch.utils.data import DataLoader
+    
+    print("=" * 60)
+    print("EXAMPLE 1: Basic Dataset Usage")
+    print("=" * 60)
+    
+    # Create dataset with both queries and database
+    dataset = VPRDataset('data')
+    print(f"Total images in dataset: {len(dataset)}")
+    print(f"Query images: {len(dataset.get_query_filenames())}")
+    print(f"Database images: {len(dataset.get_database_filenames())}")
+    print()
+    
+    # Get a single image
+    img, filename, is_query = dataset[0]
+    print(f"First image:")
+    print(f"  Filename: {filename}")
+    print(f"  Is query: {is_query}")
+    print(f"  Image shape: {img.shape}")
+    print()
+    
+    print("=" * 60)
+    print("EXAMPLE 2: Ground Truth Lookup")
+    print("=" * 60)
+    
+    # Get a query filename
+    query_files = dataset.get_query_filenames()
+    query_file = query_files[0]
+    
+    print(f"Query: {query_file}")
+    
+    # Get ground truth matches
+    matches = dataset.gt(query_file)
+    print(f"Ground truth matches ({len(matches)} images):")
+    for match in matches:
+        print(f"  - {match}")
+    print()
+    
+    print("=" * 60)
+    print("EXAMPLE 3: Create Separate Query and Database Datasets")
+    print("=" * 60)
+    
+    # Create database-only dataset
+    db_dataset = VPRDataset('data', include_queries=False, include_database=True)
+    print(f"Database-only dataset size: {len(db_dataset)}")
+    
+    # Create query-only dataset
+    query_dataset = VPRDataset('data', include_queries=True, include_database=False)
+    print(f"Query-only dataset size: {len(query_dataset)}")
+    print()
+    
+    print("=" * 60)
+    print("EXAMPLE 4: Using with DataLoader")
+    print("=" * 60)
+    
+    # Create dataloader
+    dataloader = DataLoader(dataset, batch_size=4, shuffle=False)
+    
+    # Get a batch
+    batch_imgs, batch_filenames, batch_is_query = next(iter(dataloader))
+    print(f"Batch shape: {batch_imgs.shape}")
+    print(f"Batch filenames: {batch_filenames}")
+    print(f"Batch is_query flags: {batch_is_query}")
+    print()
+    
+    print("=" * 60)
+    print("EXAMPLE 5: Get Item Info by Filename")
+    print("=" * 60)
+    
+    item_info = dataset.get_item_by_filename(query_file)
+    print(f"Full info for {query_file}:")
+    for key, value in item_info.items():
+        if key != 'path':  # Skip path for cleaner output
+            print(f"  {key}: {value}")
+    print()
+    
+    print("=" * 60)
+    print("EXAMPLE 6: Typical VPR Workflow")
+    print("=" * 60)
+    
+    print("Typical usage pattern:")
+    print("""
+    # 1. Create separate datasets
+    db_dataset = VPRDataset('data', include_queries=False)
+    query_dataset = VPRDataset('data', include_database=False)
+    
+    # 2. Extract features for all database images
+    db_features = []
+    db_filenames = []
+    for img, filename, _ in db_dataset:
+        feat = model(img.unsqueeze(0))  # Your VPR model
+        db_features.append(feat)
+        db_filenames.append(filename)
+    
+    # 3. For each query, find matches
+    for img, query_filename, _ in query_dataset:
+        # Extract query features
+        query_feat = model(img.unsqueeze(0))
+        
+        # Compute similarities with database
+        similarities = compute_similarity(query_feat, db_features)
+        
+        # Get top-K predictions
+        top_k_indices = similarities.argsort()[::-1][:10]
+        predicted_files = [db_filenames[i] for i in top_k_indices]
+        
+        # Get ground truth
+        gt_files = query_dataset.gt(query_filename)
+        
+        # Evaluate: check if any gt_files are in predicted_files
+        recall_at_10 = any(gt in predicted_files for gt in gt_files)
+    """)

model.py

@@ -0,0 +1,62 @@
+import torch 
+import torch.nn as nn
+from dataset import VPRDataset
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+
+def load_model(model_name: str, device: str = "auto") -> nn.Module: 
+    """Load a pre-trained VPR model.
+    
+    Args:
+        model_name: Name of the model to load (currently supports "eigenplaces")
+        device: Device to load model on ("auto", "cpu", "cuda")
+    
+    Returns:
+        Loaded model in evaluation mode
+    """
+    if device == "auto":
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    if model_name.lower() == "eigenplaces": 
+        model = torch.hub.load("gmberton/eigenplaces", "get_trained_model", backbone="ResNet50", fc_output_dim=2048)
+        setattr(model, "descriptor_dim", 2048)
+    else: 
+        raise ValueError(f"Model {model_name} not found")
+    
+    model = model.to(device)
+    model.eval()
+    return model
+
+
+def compute_descriptors(model: nn.Module, dataset: VPRDataset, batch_size: int = 32, device: str = "auto") -> torch.Tensor:
+    """Compute descriptors for all images in the dataset.
+    
+    Args:
+        model: Pre-trained VPR model
+        dataset: VPRDataset containing images
+        batch_size: Batch size for processing
+        device: Device to run inference on
+    
+    Returns:
+        Tensor of shape (len(dataset), descriptor_dim) containing all descriptors
+    """
+    if device == "auto":
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    model.eval()
+    descriptors = torch.zeros(len(dataset), model.descriptor_dim, device="cpu")
+    
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=2)
+    
+    with torch.no_grad():
+        start_idx = 0
+        for batch_images, batch_filenames, batch_is_query in tqdm(dataloader, desc="Computing descriptors"):
+            batch_images = batch_images.to(device)
+            batch_descriptors = model(batch_images)
+            batch_descriptors = batch_descriptors.cpu()
+            end_idx = start_idx + batch_descriptors.size(0)
+            descriptors[start_idx:end_idx] = batch_descriptors
+            start_idx = end_idx
+
+    return descriptors.cpu()  # Return on CPU for easier handling

requirements.txt

@@ -0,0 +1,6 @@
+torch>=1.9.0
+torchvision>=0.10.0
+gradio>=3.0.0
+Pillow>=8.0.0
+numpy>=1.21.0
+tqdm>=4.60.0

scripts/sample_data.py

@@ -0,0 +1,158 @@
+"""
+Sample database and query images from GSV-Cities dataset for VPR demo.
+Creates data/database/ and data/query/ folders with images.
+Ground truth is encoded in the filename: placeID_imageID.jpg
+"""
+
+import pandas as pd
+from pathlib import Path
+from PIL import Image
+import shutil
+import random
+import json
+from tqdm import tqdm
+
+# Configuration
+BASE_PATH = '/Users/olivergrainge/datasets/gsv-cities'
+OUTPUT_PATH = 'data'
+NUM_PLACES = 5  # Number of unique places to sample
+DB_IMAGES_PER_PLACE = 4  # Images per place for database
+QUERY_IMAGES_PER_PLACE = 1  # Images per place for queries
+CITIES = ['London', 'Boston', 'Melbourne']  # Cities to sample from
+MIN_IMAGES_PER_PLACE = 5  # Minimum images a place must have
+
+def load_dataframes(base_path, cities):
+    """Load and combine dataframes from multiple cities."""
+    dfs = []
+    for i, city in enumerate(cities):
+        df_path = Path(base_path) / 'Dataframes' / f'{city}.csv'
+        if not df_path.exists():
+            print(f"Warning: {df_path} not found, skipping {city}")
+            continue
+        
+        df = pd.read_csv(df_path)
+        # Add prefix to place_id to distinguish between cities
+        df['place_id'] = df['place_id'] + (i * 10**5)
+        df['city_name'] = city
+        dfs.append(df)
+    
+    if not dfs:
+        raise FileNotFoundError("No valid city dataframes found!")
+    
+    return pd.concat(dfs, ignore_index=True)
+
+def get_img_path(base_path, row):
+    """Construct the full image path from a dataframe row."""
+    city = row['city_id']
+    pl_id = row['place_id'] % 10**5
+    pl_id = str(pl_id).zfill(7)
+    panoid = row['panoid']
+    year = str(row['year']).zfill(4)
+    month = str(row['month']).zfill(2)
+    northdeg = str(row['northdeg']).zfill(3)
+    lat, lon = str(row['lat']), str(row['lon'])
+    
+    img_name = f"{city}_{pl_id}_{year}_{month}_{northdeg}_{lat}_{lon}_{panoid}.jpg"
+    return Path(base_path) / 'Images' / city / img_name
+
+def sample_and_copy_images():
+    """Main function to sample and organize images."""
+    
+    # Create output directories
+    db_path = Path(OUTPUT_PATH) / 'database'
+    query_path = Path(OUTPUT_PATH) / 'query'
+    db_path.mkdir(parents=True, exist_ok=True)
+    query_path.mkdir(parents=True, exist_ok=True)
+    
+    print("Loading dataframes...")
+    df = load_dataframes(BASE_PATH, CITIES)
+    
+    # Filter places with minimum number of images
+    place_counts = df.groupby('place_id').size()
+    valid_places = place_counts[place_counts >= (DB_IMAGES_PER_PLACE + QUERY_IMAGES_PER_PLACE)].index
+    df = df[df['place_id'].isin(valid_places)]
+    
+    print(f"Found {len(valid_places)} valid places")
+    
+    # Sample N random places
+    sampled_places = random.sample(list(valid_places), min(NUM_PLACES, len(valid_places)))
+    
+    print(f"Sampling {len(sampled_places)} places...")
+    
+    # Ground truth structure
+    ground_truth = {
+        'database': [],
+        'query': [],
+        'place_mapping': {}
+    }
+    
+    db_count = 0
+    query_count = 0
+    
+    for place_id in tqdm(sampled_places, desc="Processing places"):
+        place_images = df[df['place_id'] == place_id]
+        
+        # Sample images for this place
+        sampled = place_images.sample(n=min(DB_IMAGES_PER_PLACE + QUERY_IMAGES_PER_PLACE, len(place_images)))
+        
+        # Split into database and query
+        db_images = sampled.iloc[:DB_IMAGES_PER_PLACE]
+        query_images = sampled.iloc[DB_IMAGES_PER_PLACE:DB_IMAGES_PER_PLACE + QUERY_IMAGES_PER_PLACE]
+        
+        # Copy database images
+        for idx, (_, row) in enumerate(db_images.iterrows()):
+            src_path = get_img_path(BASE_PATH, row)
+            if not src_path.exists():
+                print(f"Warning: {src_path} not found, skipping")
+                continue
+            
+            # New filename: placeID_dbXXXX.jpg
+            dst_filename = f"place{str(place_id).zfill(8)}_db{str(db_count).zfill(4)}.jpg"
+            dst_path = db_path / dst_filename
+            
+            shutil.copy2(src_path, dst_path)
+            ground_truth['database'].append({
+                'filename': dst_filename,
+                'place_id': int(place_id),
+                'city': row['city_name'],
+                'lat': float(row['lat']),
+                'lon': float(row['lon'])
+            })
+            db_count += 1
+        
+        # Copy query images
+        for idx, (_, row) in enumerate(query_images.iterrows()):
+            src_path = get_img_path(BASE_PATH, row)
+            if not src_path.exists():
+                print(f"Warning: {src_path} not found, skipping")
+                continue
+            
+            # New filename: placeID_qXXXX.jpg
+            dst_filename = f"place{str(place_id).zfill(8)}_q{str(query_count).zfill(4)}.jpg"
+            dst_path = query_path / dst_filename
+            
+            shutil.copy2(src_path, dst_path)
+            ground_truth['query'].append({
+                'filename': dst_filename,
+                'place_id': int(place_id),
+                'city': row['city_name'],
+                'lat': float(row['lat']),
+                'lon': float(row['lon'])
+            })
+            query_count += 1
+    
+    # Save ground truth to JSON
+    gt_path = Path(OUTPUT_PATH) / 'ground_truth.json'
+    with open(gt_path, 'w') as f:
+        json.dump(ground_truth, f, indent=2)
+    
+    print(f"\n✓ Successfully created dataset!")
+    print(f"  Database images: {db_count} (in {db_path})")
+    print(f"  Query images: {query_count} (in {query_path})")
+    print(f"  Ground truth: {gt_path}")
+    print(f"\nGround truth structure:")
+    print(f"  - Filenames contain place_id: place########_db####.jpg or place########_q####.jpg")
+    print(f"  - JSON file contains detailed metadata including GPS coordinates")
+
+if __name__ == "__main__":
+    sample_and_copy_images()