geoguessr_bot/commands/kaggle_submission_command.py

import os
from dataclasses import dataclass

import pandas as pd

from geoguessr_bot.commands import AbstractCommand
from geoguessr_bot.guessr import AbstractGuessr


@dataclass
class KaggleSubmissionCommand(AbstractCommand):
    """Submit a prediction to Kaggle
    """
    image_folder_path: str
    output_path: str
    guessr: AbstractGuessr

    def run(self) -> None:
        images_ids, latitudes, longitudes = [], [], []
        for image_name in os.listdir(self.image_folder_path):
            image_path = os.path.join(self.image_folder_path, image_name)
            coordinate = self.guessr.guess_from_path(image_path)
            images_ids.append(image_name.split(".")[0])
            latitudes.append(coordinate.latitude)
            longitudes.append(coordinate.longitude)
        pd.DataFrame(dict(
            image_id=images_ids,
            latitude=latitudes,
            longitude=longitudes,
        )).to_csv(self.output_path, index=False)