classifier.py

import faiss
import numpy as np
import pickle
from glob import glob
import clip
from PIL import Image
import torch
import pycountry
import shutil
import os
import json
import matplotlib.pyplot as plt
from GeoReferencing import convert_pixel_to_latlon, haversine

# ==== Path Constants ====
GROUND_ALT = 233.591  # TODO: Set appropriately
calibration_result = np.load("./GeoReferencing/calibration_result.npz")
K = calibration_result["K"]
dist_coeffs = calibration_result["dist_coeffs"]
VIDEO_PATH = "./data/raw/GX011072.mp4"  # "./agent/input_dir/vid.mp4"

if VIDEO_PATH == "./data/raw/GX011072.mp4":
    if input("YOU ARE USING THE DEBUGGING VIDEO PATH, YOU SURE YOU WANT TO CONTINUE? (y/n) ") == "n":
        exit(0)

if input("did you set ground alt? (y/n) ") == "n":
    exit(0)

DEVICE = "cpu"
TRASH_CLASSIFICATION_DIR = "agent/trash/classification_matches"
OUTPUT_CLASSIFICATION_DIR = "agent/output/classification_matches"

REFERENCE_INDEX_PATH = "FlagClassifier/reference_data/index.faiss"
REFERENCE_EMBEDDINGS_PATH = "FlagClassifier/reference_data/embeddings.npy"
REFERENCE_METADATA_PATH = "FlagClassifier/reference_data/metadata.pkl"

INPUT_IMAGES_PATH = "agent/input_dir/target/*"
REFERENCE_COUNTRY_PATH = "FlagClassifier/reference/country"
REFERENCE_INSTITUTION_PATH = "FlagClassifier/reference/institution"
OUTPUT_MATCHES_PATH = "agent/output/classification_matches"

# ==== Directory Management ====
if os.path.exists(TRASH_CLASSIFICATION_DIR):
    shutil.rmtree(TRASH_CLASSIFICATION_DIR)

if os.path.exists(OUTPUT_CLASSIFICATION_DIR):
    os.makedirs(os.path.dirname(TRASH_CLASSIFICATION_DIR), exist_ok=True)
    shutil.move(OUTPUT_CLASSIFICATION_DIR, TRASH_CLASSIFICATION_DIR)

os.makedirs(OUTPUT_CLASSIFICATION_DIR, exist_ok=True)

# ==== Model and Data Loading ====
model, preprocess = clip.load("ViT-B/16", device=DEVICE)
index = faiss.read_index(REFERENCE_INDEX_PATH)
all_emb = np.load(REFERENCE_EMBEDDINGS_PATH)
metadata = pickle.load(open(REFERENCE_METADATA_PATH, "rb"))
all_labels = metadata["labels"]
all_classes = metadata["classes"]

images_path = glob(INPUT_IMAGES_PATH)


# ==== Ensemble Suggestions ====
ensambled_suggestions = {}
for img_path in images_path:
    crop = Image.open(img_path).convert("RGB")
    crop = preprocess(crop).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        crop_emb = model.encode_image(crop).cpu().numpy().astype("float32")
    faiss.normalize_L2(crop_emb)

    D, I = index.search(crop_emb, k=25)
    suggestions = [(all_labels[i], all_classes[i], all_emb[i]) for i in I[0]]

    for i, suggestion in enumerate(suggestions):
        label = suggestion[0]
        score = 1.0 / (i + 1)
        ensambled_suggestions[label] = ensambled_suggestions.get(label, 0) + score


# ==== Output Results ====

code_to_name = {c.alpha_2.lower(): c.name for c in pycountry.countries}  # type: ignore

# Load posxy.json for pixel positions
with open("agent/output/posxy.json", "r") as f:
    posxy_data = json.load(f)


# Helper to extract frame_idx and box_idx from image filename
def extract_indices(img_path):
    # Example: frame_2370_crop_0.jpg
    base = os.path.basename(img_path)
    parts = base.replace(".jpg", "").split("_")
    frame_idx = int(parts[1])
    box_idx = int(parts[-1])
    return frame_idx, box_idx


latlon_array = []


print("\nFinal ensembled scores:")
i = 1
for label, score in sorted(ensambled_suggestions.items(), key=lambda x: x[1], reverse=True):
    display_label = f"{label}_{code_to_name[label]}" if label in code_to_name else label

    src_img_path = f"{REFERENCE_COUNTRY_PATH}/{label}.png"
    if not os.path.exists(src_img_path):
        src_img_path = f"{REFERENCE_INSTITUTION_PATH}/{label}.png"

    if os.path.exists(src_img_path):
        os.makedirs(OUTPUT_MATCHES_PATH, exist_ok=True)
        dst_img_path = f"{OUTPUT_MATCHES_PATH}/{i}_{display_label}.png"
        shutil.copy(src_img_path, dst_img_path)
    else:
        print(f"Image not found for {label}, skipping.")
    print(f"{score:.3f}_{display_label}.png")
    i += 1

# Lat/Lon extraction loop (unchanged)
latlon_array = []
for img_path in images_path:
    frame_idx, box_idx = extract_indices(img_path)
    posxy = next((item for item in posxy_data if item["frame_idx"] == frame_idx and item["box_idx"] == box_idx), None)
    if posxy:
        pos_x = posxy["box_center_x"]
        pos_y = posxy["box_center_y"]
        lat, lon = convert_pixel_to_latlon(VIDEO_PATH, frame_idx, pos_x, pos_y, GROUND_ALT, K, dist_coeffs)
        latlon_array.append([lat, lon])
        print(f"Lat/Lon: {lat}, {lon}")
    else:
        print(f"No posxy data for {img_path}")


if latlon_array:
    mean_lat = np.mean([lat for lat, lon in latlon_array])
    mean_lon = np.mean([lon for lat, lon in latlon_array])
    print(f"Mean Lat/Lon: [{mean_lat}, {mean_lon}]")

    # Calculate error between mean lat/lon and the correct lat/lon
    correct_latlon = [29.8189133, 30.82604093]
    if latlon_array:
        mean_error = haversine(mean_lat, mean_lon, correct_latlon[0], correct_latlon[1])
        print(f"Error between mean location and correct location: {mean_error:.2f} meters")
else:
    print("WHY THE HELL IS THERE NO LATLON_ARRAY")

# show distribution
# Plot all lat/lon points
# if latlon_array:
#     lats, lons = zip(*latlon_array)
#     plt.figure(figsize=(8, 6))
#     plt.scatter(lons, lats, c="red", marker="o")
#     plt.xlabel("Longitude")
#     plt.ylabel("Latitude")
#     plt.title("Flag Locations")
#     plt.grid(True)
#     plt.show()

# for debugging only show how far we were
# if latlon_array:
#     latlon_array.append([29.8189133, 30.82604093])  # Add the correct lat/lon for testing
#     dist_matrix = np.zeros((len(latlon_array), len(latlon_array)))
#     for i, (lat1, lon1) in enumerate(latlon_array):
#         for j, (lat2, lon2) in enumerate(latlon_array):
#             dist_matrix[i, j] = haversine(lat1, lon1, lat2, lon2)

#     print("Pairwise Haversine Distance Matrix (meters):")
#     print(dist_matrix)

#     plt.figure(figsize=(8, 6))
#     plt.imshow(dist_matrix, cmap="viridis")
#     plt.colorbar(label="Distance (meters)")
#     plt.title("Pairwise Haversine Distances Between Flag Locations")
#     plt.xlabel("Point Index")
#     plt.ylabel("Point Index")
#     plt.show()