Spaces:
Runtime error
Runtime error
File size: 2,629 Bytes
1791df2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
from collections import Counter
from dataclasses import dataclass
import numpy as np
from sklearn.cluster import DBSCAN
from sklearn.metrics.pairwise import haversine_distances
from PIL import Image
import pandas as pd
from geoguessr_bot.guessr import AbstractGuessr
from geoguessr_bot.interfaces import Coordinate
from geoguessr_bot.retriever import AbstractImageEmbedder
from geoguessr_bot.retriever import Retriever
@dataclass
class AverageNeighborsEmbedderGuessr(AbstractGuessr):
"""Guesses a coordinate using an Embedder and a retriever followed by NN.
"""
embedder: AbstractImageEmbedder
retriever: Retriever
metadata_path: str
n_neighbors: int = 1000
dbscan_eps: float = 0.05
def __post_init__(self):
"""Load metadata
"""
metadata = pd.read_csv(self.metadata_path)
self.image_to_coordinate = {
image.split("/")[-1]: Coordinate(latitude=latitude, longitude=longitude)
for image, latitude, longitude in zip(metadata["path"], metadata["latitude"], metadata["longitude"])
}
# DBSCAN will be used to take the centroid of the biggest cluster among the N neighbors, using Haversine
self.dbscan = DBSCAN(eps=self.dbscan_eps, metric=haversine_distances)
def guess(self, image: Image) -> Coordinate:
"""Guess a coordinate from an image
"""
# Embed image
image = Image.fromarray(image)
image_embedding = self.embedder.embed(image)[None, :]
# Retrieve nearest neighbors
nearest_neighbors, distances = self.retriever.retrieve(image_embedding, self.n_neighbors)
nearest_neighbors = nearest_neighbors[0]
distances = distances[0]
# Get coordinates of neighbors
neighbors_coordinates = [self.image_to_coordinate[nn].to_radians() for nn in nearest_neighbors]
neighbors_coordinates = np.array([[nn.latitude, nn.longitude] for nn in neighbors_coordinates])
# Use DBSCAN to find the biggest cluster and potentially remove outliers
clustering = self.dbscan.fit(neighbors_coordinates)
labels = clustering.labels_
biggest_cluster = max(Counter(labels))
neighbors_coordinates = neighbors_coordinates[labels == biggest_cluster]
distances = distances[labels == biggest_cluster]
# Guess coordinate as the closest image among the cluster regarding retrieving distance
guess_coordinate = neighbors_coordinates[np.argmin(distances)]
guess_coordinate = Coordinate.from_radians(guess_coordinate[0], guess_coordinate[1])
return guess_coordinate
|