import asyncio import logging import os import shutil from pathlib import Path import time import numpy as np import torch import yaml from google import genai from google.genai import types from pydantic import ValidationError from tqdm.asyncio import tqdm as atqdm from .data_processor import DataProcessor from .g3.G3 import G3 from .prompt import ( Evidence, GPSPrediction, LocationPrediction, diversification_prompt, location_prompt, verification_prompt, ) from .utils import ( calculate_similarity_scores, extract_and_parse_json, get_gps_from_location, handle_async_api_call_with_retry, image_to_base64, ) logger = logging.getLogger("uvicorn.error") class G3BatchPredictor: """ Batch prediction class for processing all images and videos in a directory. This class: 1. Preprocesses all images and videos in a directory. 2. Extracts keyframes from videos and combines them with images. 3. Passes all keyframes and images to the Gemini model for prediction. """ def __init__( self, device: str = "cuda", input_dir: str = "data/input_data", prompt_dir: str = "data/prompt_data", cache_dir: str = "data/cache", index_path: str = "data/index/G3.index", hparams_path: str = "g3/hparams.yaml", database_csv_path: str = "data/dataset/mp16/MP16_Pro_filtered.csv", checkpoint_path: str = "data/checkpoints/mercator_finetune_weight.pth", ): """ Initialize the BatchKeyframePredictor. Args: checkpoint_path (str): Path to G3 model checkpoint device (str): Device to run model on ("cuda" or "cpu") index_path (str): Path to FAISS index for RAG (required) """ self.device = torch.device(device) self.base_path = Path(__file__).parent self.checkpoint_path = self.base_path / checkpoint_path self.input_dir = self.base_path / input_dir self.prompt_dir = self.base_path / prompt_dir self.cache_dir = self.base_path / cache_dir self.image_dir = self.prompt_dir / "images" self.audio_dir = self.prompt_dir / "audio" os.makedirs(self.input_dir, exist_ok=True) os.makedirs(self.prompt_dir, exist_ok=True) os.makedirs(self.cache_dir, exist_ok=True) os.makedirs(self.image_dir, exist_ok=True) os.makedirs(self.audio_dir, exist_ok=True) # Initialize G3 model hparams = yaml.safe_load(open(self.base_path / hparams_path, "r")) pe = "projection_mercator" nn = "rffmlp" self.model = G3( device=device, positional_encoding_type=pe, neural_network_type=nn, hparams=hparams[f"{pe}_{nn}"], ) self.__load_checkpoint() self.data_processor = DataProcessor( model=self.model, input_dir=self.input_dir, prompt_dir=self.prompt_dir, cache_dir=self.cache_dir, image_dir=self.image_dir, audio_dir=self.audio_dir, index_path=self.base_path / index_path, database_csv_path=self.base_path / database_csv_path, device=self.device, ) self.image_extension = { ".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp", } self.video_extension = { ".mp4", ".avi", ".mov", ".mkv", } def __load_checkpoint(self): """ Load the G3 model checkpoint. """ if not os.path.exists(self.checkpoint_path): raise FileNotFoundError( f"Checkpoint file not found: {self.checkpoint_path}" ) self.model.load_state_dict( torch.load(self.checkpoint_path, map_location=self.device) ) self.model.to(self.device) self.model.eval() logger.info( f"āœ… Successfully loaded G3 model checkpoint from: {self.checkpoint_path}" ) async def llm_predict( self, model_name: str = "gemini-2.5-pro", n_search: int | None = None, n_coords: int | None = None, image_prediction: bool = True, text_prediction: bool = True, ) -> LocationPrediction: """ Generate a prediction using the Gemini LLM with Pydantic structured output. Args: model_name: LLM model name to use n_search: Number of search results to include n_coords: Number of coordinates to include image_prediction: Whether to use images in prediction text_prediction: Whether to use text in prediction Returns: dict: Parsed prediction response """ prompt = diversification_prompt( prompt_dir=str(self.prompt_dir), n_coords=n_coords, n_search=n_search, image_prediction=image_prediction, text_prediction=text_prediction, ) images = [] if image_prediction: image_dir = self.image_dir if not image_dir.exists(): raise ValueError(f"Image directory does not exist: {image_dir}") for image_file in image_dir.glob("*.jpg"): with open(image_file, "rb") as f: image = types.Part.from_bytes(data=f.read(), mime_type="image/jpeg") images.append(image) client = genai.Client(api_key=os.environ["GOOGLE_CLOUD_API_KEY"]) async def api_call(): loop = asyncio.get_event_loop() response = await loop.run_in_executor( None, lambda: client.models.generate_content( model=model_name, contents=[*images, prompt], config=types.GenerateContentConfig( tools=[ types.Tool(url_context=types.UrlContext()), ], temperature=0.1, top_p=0.95, ), ), ) raw_text = response.text.strip() if response.text is not None else "" parsed_json = extract_and_parse_json(raw_text) try: validated = LocationPrediction.model_validate(parsed_json) return validated except (ValidationError, ValueError): raise ValueError("Empty or invalid LLM response") return await handle_async_api_call_with_retry( api_call, fallback_result=LocationPrediction( latitude=0.0, longitude=0.0, location="", evidence=[] ), error_context=f"LLM prediction with {model_name}", ) async def diversification_predict( self, model_name: str = "gemini-2.5-flash", image_prediction: bool = True, text_prediction: bool = True, ) -> LocationPrediction: """ Diversification prediction without preprocessing (assumes preprocessing already done). Runs different sample sizes in parallel for faster execution. Args: model_name (str): LLM model name to use image_prediction (bool): Whether to use images in prediction text_prediction (bool): Whether to use text in prediction Returns: dict: Best prediction with latitude, longitude, location, reason, and metadata """ # Function to try a specific sample size with retry logic async def try_sample_size(num_sample): while True: prediction = await self.llm_predict( model_name=model_name, n_search=num_sample, n_coords=num_sample, image_prediction=image_prediction, text_prediction=text_prediction, ) if prediction: coords = (prediction.latitude, prediction.longitude) return (num_sample, coords, prediction) else: logger.info( f"Invalid or empty prediction format with {num_sample} samples, retrying..." ) # Run all sample sizes in parallel num_samples = [10, 15, 20] logger.info( f"šŸš€ Running {len(num_samples)} sample sizes in parallel: {num_samples}" ) tasks = [try_sample_size(num_sample) for num_sample in num_samples] class LW: def write(self, msg: str) -> int: logger.info(msg) return len(msg) def flush(self): pass results = await atqdm.gather( *tasks, desc="šŸ”„ Running diversification predictions", file=LW(), ) # Build predictions dictionary from parallel results predictions_dict = {} for num_sample, coords, prediction in results: predictions_dict[coords] = prediction logger.info(f"āœ… Collected prediction with {num_sample} samples: {coords}") # Convert predictions to coordinate list for similarity scoring predicted_coords = list(predictions_dict.keys()) logger.info(f"Predicted coordinates: {predicted_coords}") if not predicted_coords: raise ValueError("No valid predictions obtained from any sample size") # Calculate similarity scores avg_similarities = calculate_similarity_scores( model=self.model, device=self.device, predicted_coords=predicted_coords, image_dir=self.image_dir, ) # Find best prediction best_idx = np.argmax(avg_similarities) best_coords = predicted_coords[best_idx] best_prediction = predictions_dict[best_coords] logger.info(f"šŸŽÆ Best prediction selected: {best_coords}") logger.info(f" Similarity scores: {avg_similarities}") logger.info(f" Best index: {best_idx}") # print(json.dumps(best_prediction, indent=2)) # Commented out verbose output return best_prediction async def location_predict( self, model_name: str = "gemini-2.5-flash", location: str = "specified location", ) -> GPSPrediction: """ Generate a location-based prediction using the Gemini LLM with centralized retry logic. Args: model_name (str): LLM model name to use location (str): Location to use in the prompt Returns: dict: Parsed JSON prediction response """ if not location: raise ValueError("Location must be specified for location-based prediction") lat, lon = get_gps_from_location(location) if lat is not None and lon is not None: logger.info( f"Using GPS coordinates for location '{location}': ({lat}, {lon})" ) return GPSPrediction( latitude=lat, longitude=lon, analysis="", references=[] ) else: prompt = location_prompt(location) client = genai.Client(api_key=os.environ["GOOGLE_CLOUD_API_KEY"]) async def api_call(): # Run the synchronous API call in a thread executor to make it truly async loop = asyncio.get_event_loop() response = await loop.run_in_executor( None, lambda: client.models.generate_content( model=model_name, contents=[prompt], config=types.GenerateContentConfig( tools=[ types.Tool(google_search=types.GoogleSearch()), ], temperature=0.1, top_p=0.95, ), ), ) raw_text = response.text.strip() if response.text is not None else "" parsed_json = extract_and_parse_json(raw_text) try: validated = GPSPrediction.model_validate(parsed_json) return validated except (ValidationError, ValueError): raise ValueError("Empty or invalid LLM response") return await handle_async_api_call_with_retry( api_call, fallback_result=GPSPrediction( latitude=0.0, longitude=0.0, analysis="", references=[] ), error_context=f"Location prediction for '{location}' with {model_name}", ) async def verification_predict( self, prediction: LocationPrediction, model_name: str = "gemini-2.5-flash", image_prediction: bool = True, text_prediction: bool = True, ) -> LocationPrediction: """ Generate verification prediction based on the provided prediction. Args: prediction (dict): Prediction dictionary with latitude, longitude, location, reason, and metadata model_name (str): LLM model name to use for verification Returns: dict: Verification prediction with latitude, longitude, location, reason, and evidence """ # Prepare verification data (now async) satellite_image_id = await self.data_processor.prepare_location_images( prediction=prediction.model_dump(), image_prediction=image_prediction, text_prediction=text_prediction, ) image_dir = self.image_dir images = [] if image_prediction: if not image_dir.exists(): raise ValueError(f"Image directory does not exist: {image_dir}") for image_file in image_dir.glob("*.jpg"): with open(image_file, "rb") as f: image = types.Part.from_bytes(data=f.read(), mime_type="image/jpeg") images.append(image) # Prepare verification prompt prompt = verification_prompt( satellite_image_id=satellite_image_id, prediction=prediction.model_dump(), prompt_dir=str(self.prompt_dir), image_prediction=image_prediction, text_prediction=text_prediction, ) client = genai.Client(api_key=os.environ["GOOGLE_CLOUD_API_KEY"]) async def api_call(): # Run the synchronous API call in a thread executor to make it truly async loop = asyncio.get_event_loop() response = await loop.run_in_executor( None, lambda: client.models.generate_content( model=model_name, contents=[*images, prompt], config=types.GenerateContentConfig( tools=[ types.Tool(url_context=types.UrlContext()), ], temperature=0.1, top_p=0.95, ), ), ) raw_text = response.text.strip() if response.text is not None else "" parsed_json = extract_and_parse_json(raw_text) try: validated = LocationPrediction.model_validate(parsed_json) return validated except (ValidationError, ValueError): raise ValueError("Empty or invalid LLM response") return await handle_async_api_call_with_retry( api_call, fallback_result=LocationPrediction( latitude=0.0, longitude=0.0, location="", evidence=[] ), error_context=f"Verification prediction with {model_name}", ) async def predict( self, model_name: str = "gemini-2.5-flash", image_prediction: bool = True, text_prediction: bool = True, ) -> LocationPrediction: """ Complete prediction pipeline without preprocessing (assumes preprocessing already done). Used for parallel execution where preprocessing is done once beforehand. All major steps run in parallel for maximum speed. Args: model_name (str): LLM model name to use image_prediction (bool): Whether to use images in prediction text_prediction (bool): Whether to use text in prediction Returns: dict: Final prediction with latitude, longitude, location, reason, and evidence """ logger.info( f"šŸš€ Starting multi-modal prediction pipeline with model: {model_name}" ) await self.data_processor.preprocess_input_data() # Step 1: Run diversification prediction (this is already parallel internally) logger.info( f"\nšŸ”„ Running diversification prediction for Image={image_prediction}, Text={text_prediction}..." ) diversification_result = await self.diversification_predict( model_name=model_name, image_prediction=image_prediction, text_prediction=text_prediction, ) # Step 2: Run location prediction location_prediction = await self.location_predict( model_name=model_name, location=diversification_result.location ) logger.info("āœ… Location prediction completed:") # Step 3: Update coordinates and evidence from location prediction result = diversification_result.model_copy() result.longitude = location_prediction.longitude result.latitude = location_prediction.latitude # Step 4: Normalize and append location evidence if location_prediction.analysis and location_prediction.references: location_evidence = Evidence( analysis=location_prediction.analysis, references=location_prediction.references, ) else: location_evidence = Evidence( analysis="No specific location analysis provided.", references=[], ) # Append to result evidence result.evidence.append(location_evidence) # Step 5: Run verification prediction logger.info( f"\nšŸ”„ Running verification prediction for Image={image_prediction}, Text={text_prediction}..." ) result = await self.verification_predict( prediction=result, model_name=model_name, image_prediction=image_prediction, text_prediction=text_prediction, ) logger.info( f"\nšŸŽÆ Final prediction for Image={image_prediction}, Text={text_prediction}:" ) # print(json.dumps(result, indent=2)) # Commented out verbose output return result def get_response(self, prediction: LocationPrediction) -> LocationPrediction: """ Convert image references in the prediction to base64 strings. """ for evidence in prediction.evidence: for i, ref in enumerate(evidence.references): if ref.startswith("image"): evidence.references[i] = image_to_base64(self.image_dir / ref) return prediction def get_transcript(self) -> str: """ Get the transcript from the transcript files in the audio directory. """ transcript = "" for transcript_file in self.audio_dir.glob("*.txt"): with open(transcript_file, "r", encoding="utf-8") as f: logger.info(f"Reading transcript from {transcript_file.name}") transcript_data = f.read().strip() if transcript_data: transcript += f"Transcript for {transcript_file.name}\n" transcript += transcript_data return transcript def clear_directories(self): """ Clear the input and prompt directories. """ delete_dirs = [self.input_dir, self.prompt_dir] for dir_path in delete_dirs: if os.path.exists(dir_path): shutil.rmtree(dir_path) logger.info(f"Deleted folder: {dir_path}") else: logger.info(f"Folder does not exist: {dir_path}")