Spaces:

SalHargis
/

UniqueAnimalReIdentification

Sleeping

Update app.py

8c24824 verified 9 months ago

10 kB

import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import os
import numpy as np
import matplotlib.pyplot as plt
import random

# LightGlue Imports
from lightglue import LightGlue, ALIKED
from lightglue.utils import load_image, rbd

# Configuration
WEIGHTS_PATH = "MiewID_ArcFace_FineTun.pth"
GALLERY_FILE = "mini_gallery.pt"
TEST_QUERIES_DIR = "test_queries"
IMG_SIZE = 384
EMBEDDING_DIM = 512
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MAX_VIZ_LINES = 50

# Lets load the models
print("Loading Models...")
model = models.resnet50(weights=None)
model.fc = nn.Linear(model.fc.in_features, EMBEDDING_DIM)
model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=DEVICE), strict=False)
model.to(DEVICE).eval()

extractor = ALIKED(max_num_keypoints=1024, detection_threshold=0.2).eval().to(DEVICE)
matcher = LightGlue(features='aliked').eval().to(DEVICE)

# load the gallery of images
if os.path.exists(GALLERY_FILE):
    data = torch.load(GALLERY_FILE, map_location=DEVICE)
    g_embeddings = data["embeddings"].to(DEVICE)
    g_paths = data["paths"]
    g_labels = data["labels"]
    g_species = data["species"]
else:
    raise FileNotFoundError("Gallery file missing!")

# ground truth lookup
GT_LOOKUP = {}
if os.path.exists(TEST_QUERIES_DIR):
    for f in os.listdir(TEST_QUERIES_DIR):
        if f.lower().endswith(('.jpg', '.png', '.jpeg')):
            full_path = os.path.join(TEST_QUERIES_DIR, f)
            try:
                f_size = os.path.getsize(full_path)
                parts = f.split("_")
                if len(parts) >= 2:
                    GT_LOOKUP[f_size] = (parts[0], parts[1])
            except: pass

transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def create_match_visualization(img_path1, img_path2, kpts0, kpts1, matches, output_name, max_lines=MAX_VIZ_LINES):
    if kpts0.dim() == 3: kpts0 = kpts0[0]
    if kpts1.dim() == 3: kpts1 = kpts1[0]

    # Density redction
    num_matches = len(matches)
    if num_matches > max_lines:
        indices = random.sample(range(num_matches), max_lines)
        matches_to_draw = matches[indices]
    else:
        matches_to_draw = matches

    img1 = Image.open(img_path1).convert("RGB")
    img2 = Image.open(img_path2).convert("RGB")
    
    target_h = 400
    w1, h1 = img1.size
    w2, h2 = img2.size
    scale1 = target_h / h1
    scale2 = target_h / h2
    
    img1 = img1.resize((int(w1 * scale1), target_h))
    img2 = img2.resize((int(w2 * scale2), target_h))
    
    fig, ax = plt.subplots(1, 1, figsize=(10, 5))
    ax.axis('off')
    
    img1_np = np.array(img1)
    img2_np = np.array(img2)
    h1_new, w1_new, _ = img1_np.shape
    h2_new, w2_new, _ = img2_np.shape
    
    width = w1_new + w2_new
    canvas = np.zeros((target_h, width, 3), dtype=np.uint8)
    canvas[:, :w1_new, :] = img1_np
    canvas[:, w1_new:, :] = img2_np
    
    ax.imshow(canvas)
    
    m_kpts0 = kpts0[matches_to_draw[..., 0]].cpu().numpy()
    m_kpts1 = kpts1[matches_to_draw[..., 1]].cpu().numpy()
    
    m_kpts0[:, 0] *= scale1
    m_kpts0[:, 1] *= scale1
    m_kpts1[:, 0] *= scale2
    m_kpts1[:, 1] *= scale2
    
    for (x0, y0), (x1, y1) in zip(m_kpts0, m_kpts1):
        ax.plot([x0, x1 + w1_new], [y0, y1], color="lime", linewidth=0.8, alpha=0.6)
        ax.scatter([x0, x1 + w1_new], [y0, y1], color="lime", s=3)

    plt.tight_layout()
    output_path = f"{output_name}.jpg"
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0, dpi=150)
    plt.close()
    return output_path

def predict(input_path):
    # default return values are emopty
    default_header = "Please upload an image."
    default_logs = ["", "", ""]
    default_imgs = [None, None, None]
    
    if input_path is None: 
        return default_header, default_logs[0], default_imgs[0], default_logs[1], default_imgs[1], default_logs[2], default_imgs[2]
    
    # Ground truth
    true_species, true_id = "Unknown", "Unknown"
    try:
        input_size = os.path.getsize(input_path)
        if input_size in GT_LOOKUP:
            true_species, true_id = GT_LOOKUP[input_size]
        else:
            filename = os.path.basename(input_path)
            if "_QUERY" in filename:
                parts = filename.split("_")
                true_species, true_id = parts[0], parts[1]
    except: pass

    # load Image
    input_image = Image.open(input_path).convert("RGB")

    # Coarse Search (ArcFace)
    img_t = transform(input_image).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        q_emb = torch.nn.functional.normalize(model(img_t), p=2, dim=1)
        scores = torch.mm(q_emb, g_embeddings.t())
        top_scores, top_indices = torch.topk(scores, k=min(50, len(g_paths)))

    # Filter for unique individual candidates
    unique_candidates = []
    seen_individuals = set()
    for i in range(len(top_indices[0])):
        if len(unique_candidates) >= 3: break
        idx = top_indices[0][i].item()
        score = top_scores[0][i].item()
        label = g_labels[idx]
        if label not in seen_individuals:
            seen_individuals.add(label)
            unique_candidates.append((idx, score))

    # Fine search and visualization 
    feats_q = extractor.extract(load_image(input_path).to(DEVICE))
    
    best_score = -1
    best_candidate_idx = -1
    
    # Initialize output lists (size 3)
    cand_logs = ["Waiting for data...", "Waiting for data...", "Waiting for data..."]
    cand_viz_paths = [None, None, None]

    for rank, (idx, arcface_sim) in enumerate(unique_candidates):
        path = g_paths[idx]
        label = g_labels[idx]
        species = g_species[idx]
        
        try:
            if not os.path.exists(path): continue

            feats_c = extractor.extract(load_image(path).to(DEVICE))
            with torch.no_grad():
                matches = matcher({"image0": feats_q, "image1": feats_c})
                matches = rbd(matches)
            
            geo_matches = len(matches["matches"])
            sim_percent = arcface_sim * 100
            
            # Create individual log string
            log_str = f"### Candidate {rank+1}: {species} / {label}\n"
            log_str += f"**Coarse-Search Confidence:** {sim_percent:.1f}% &nbsp;&nbsp;|&nbsp;&nbsp; **📐 Geometric Matches:** {geo_matches}"
            cand_logs[rank] = log_str
            
            viz_name = f"viz_rank_{rank}"
            viz_path = create_match_visualization(
                input_path, path, 
                feats_q['keypoints'], feats_c['keypoints'], 
                matches['matches'], viz_name
            )
            cand_viz_paths[rank] = viz_path

            if geo_matches > best_score:
                best_score = geo_matches
                best_candidate_idx = idx
                
        except Exception as e:
            cand_logs[rank] = f"Error processing candidate: {e}"

    # Final decision calculation
    CONFIDENCE_THRESHOLD = 15
    if best_candidate_idx != -1 and best_score > CONFIDENCE_THRESHOLD:
        pred_species = g_species[best_candidate_idx]
        pred_id = g_labels[best_candidate_idx]
        is_correct = (pred_id == true_id)
        
        if true_id == "Unknown": header = f"### ❓ MATCH FOUND (No Ground Truth)\n"
        elif is_correct: header = f"### ✅ CORRECT MATCH!\n"
        else: header = f"### ❌ INCORRECT MATCH\n"
        
        header += f"**Ground Truth:** {true_species} / {true_id} &nbsp;&nbsp;➡️&nbsp;&nbsp; **Prediction:** {pred_species} / {pred_id}\n"
        header += f"*(Confirmed with {best_score} geometric keypoints)*"
    else:
        header = "### ⚠️ UNKNOWN / NO MATCH\n"
        header += f"**Ground Truth:** {true_species} / {true_id}\n"
        header += f"**Prediction:** None (Best match only had {best_score} keypoints)\n"

    # Return: Header, then (Log, Img) for Cand 1, then (Log, Img) for Cand 2, etc.
    return (header, 
            cand_logs[0], cand_viz_paths[0],
            cand_logs[1], cand_viz_paths[1],
            cand_logs[2], cand_viz_paths[2])

# Setup for the user interface 
examples_list = []
if os.path.exists(TEST_QUERIES_DIR):
    examples_list = [os.path.join(TEST_QUERIES_DIR, f) for f in os.listdir(TEST_QUERIES_DIR) if f.lower().endswith(('.jpg', '.png'))]

with gr.Blocks(title="Wildlife Re-ID: Coarse-to-Fine Demo") as demo:
    gr.Markdown("# Wildlife Re-ID: Coarse-to-Fine Demo")
    gr.Markdown("Select a test image. The system finds the Top 3 UNIQUE individuals using embeddings, then verifies them using geometry.")

    with gr.Row():
        # Left Column: Input
        with gr.Column(scale=1):
            input_img = gr.Image(type="filepath", label="Test Image", height=300)
            gr.Examples(examples=examples_list, inputs=input_img, label="Test Examples", examples_per_page=4)
            submit_btn = gr.Button("Run Identification", variant="primary", size="lg")

        # Right Column: Vertical Stack of Candidates
        with gr.Column(scale=2):
            header_md = gr.Markdown(label="Final Decision")
            
            # Candidate 1 Group
            with gr.Group():
                log1 = gr.Markdown()
                # FIX: Removed show_download_button AND height
                img1 = gr.Image(label="Visualization", show_label=False)
            
            # Candidate 2 Group
            with gr.Group():
                log2 = gr.Markdown()
                img2 = gr.Image(label="Visualization", show_label=False)
                
            # Candidate 3 Group
            with gr.Group():
                log3 = gr.Markdown()
                img3 = gr.Image(label="Visualization", show_label=False)

    submit_btn.click(
        fn=predict,
        inputs=input_img,
        outputs=[header_md, log1, img1, log2, img2, log3, img3]
    )

demo.launch(allowed_paths=[TEST_QUERIES_DIR])

	import gradio as gr
	import torch
	import torch.nn as nn
	from torchvision import models, transforms
	from PIL import Image
	import os
	import numpy as np
	import matplotlib.pyplot as plt
	import random

	# LightGlue Imports
	from lightglue import LightGlue, ALIKED
	from lightglue.utils import load_image, rbd

	# Configuration
	WEIGHTS_PATH = "MiewID_ArcFace_FineTun.pth"
	GALLERY_FILE = "mini_gallery.pt"
	TEST_QUERIES_DIR = "test_queries"
	IMG_SIZE = 384
	EMBEDDING_DIM = 512
	DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	MAX_VIZ_LINES = 50

	# Lets load the models
	print("Loading Models...")
	model = models.resnet50(weights=None)
	model.fc = nn.Linear(model.fc.in_features, EMBEDDING_DIM)
	model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=DEVICE), strict=False)
	model.to(DEVICE).eval()

	extractor = ALIKED(max_num_keypoints=1024, detection_threshold=0.2).eval().to(DEVICE)
	matcher = LightGlue(features='aliked').eval().to(DEVICE)

	# load the gallery of images
	if os.path.exists(GALLERY_FILE):
	data = torch.load(GALLERY_FILE, map_location=DEVICE)
	g_embeddings = data["embeddings"].to(DEVICE)
	g_paths = data["paths"]
	g_labels = data["labels"]
	g_species = data["species"]
	else:
	raise FileNotFoundError("Gallery file missing!")

	# ground truth lookup
	GT_LOOKUP = {}
	if os.path.exists(TEST_QUERIES_DIR):
	for f in os.listdir(TEST_QUERIES_DIR):
	if f.lower().endswith(('.jpg', '.png', '.jpeg')):
	full_path = os.path.join(TEST_QUERIES_DIR, f)
	try:
	f_size = os.path.getsize(full_path)
	parts = f.split("_")
	if len(parts) >= 2:
	GT_LOOKUP[f_size] = (parts[0], parts[1])
	except: pass

	transform = transforms.Compose([
	transforms.Resize((IMG_SIZE, IMG_SIZE)),
	transforms.ToTensor(),
	transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
	])

	def create_match_visualization(img_path1, img_path2, kpts0, kpts1, matches, output_name, max_lines=MAX_VIZ_LINES):
	if kpts0.dim() == 3: kpts0 = kpts0[0]
	if kpts1.dim() == 3: kpts1 = kpts1[0]

	# Density redction
	num_matches = len(matches)
	if num_matches > max_lines:
	indices = random.sample(range(num_matches), max_lines)
	matches_to_draw = matches[indices]
	else:
	matches_to_draw = matches

	img1 = Image.open(img_path1).convert("RGB")
	img2 = Image.open(img_path2).convert("RGB")

	target_h = 400
	w1, h1 = img1.size
	w2, h2 = img2.size
	scale1 = target_h / h1
	scale2 = target_h / h2

	img1 = img1.resize((int(w1 * scale1), target_h))
	img2 = img2.resize((int(w2 * scale2), target_h))

	fig, ax = plt.subplots(1, 1, figsize=(10, 5))
	ax.axis('off')

	img1_np = np.array(img1)
	img2_np = np.array(img2)
	h1_new, w1_new, _ = img1_np.shape
	h2_new, w2_new, _ = img2_np.shape

	width = w1_new + w2_new
	canvas = np.zeros((target_h, width, 3), dtype=np.uint8)
	canvas[:, :w1_new, :] = img1_np
	canvas[:, w1_new:, :] = img2_np

	ax.imshow(canvas)

	m_kpts0 = kpts0[matches_to_draw[..., 0]].cpu().numpy()
	m_kpts1 = kpts1[matches_to_draw[..., 1]].cpu().numpy()

	m_kpts0[:, 0] *= scale1
	m_kpts0[:, 1] *= scale1
	m_kpts1[:, 0] *= scale2
	m_kpts1[:, 1] *= scale2

	for (x0, y0), (x1, y1) in zip(m_kpts0, m_kpts1):
	ax.plot([x0, x1 + w1_new], [y0, y1], color="lime", linewidth=0.8, alpha=0.6)
	ax.scatter([x0, x1 + w1_new], [y0, y1], color="lime", s=3)

	plt.tight_layout()
	output_path = f"{output_name}.jpg"
	plt.savefig(output_path, bbox_inches='tight', pad_inches=0, dpi=150)
	plt.close()
	return output_path

	def predict(input_path):
	# default return values are emopty
	default_header = "Please upload an image."
	default_logs = ["", "", ""]
	default_imgs = [None, None, None]

	if input_path is None:
	return default_header, default_logs[0], default_imgs[0], default_logs[1], default_imgs[1], default_logs[2], default_imgs[2]

	# Ground truth
	true_species, true_id = "Unknown", "Unknown"
	try:
	input_size = os.path.getsize(input_path)
	if input_size in GT_LOOKUP:
	true_species, true_id = GT_LOOKUP[input_size]
	else:
	filename = os.path.basename(input_path)
	if "_QUERY" in filename:
	parts = filename.split("_")
	true_species, true_id = parts[0], parts[1]
	except: pass

	# load Image
	input_image = Image.open(input_path).convert("RGB")

	# Coarse Search (ArcFace)
	img_t = transform(input_image).unsqueeze(0).to(DEVICE)
	with torch.no_grad():
	q_emb = torch.nn.functional.normalize(model(img_t), p=2, dim=1)
	scores = torch.mm(q_emb, g_embeddings.t())
	top_scores, top_indices = torch.topk(scores, k=min(50, len(g_paths)))

	# Filter for unique individual candidates
	unique_candidates = []
	seen_individuals = set()
	for i in range(len(top_indices[0])):
	if len(unique_candidates) >= 3: break
	idx = top_indices[0][i].item()
	score = top_scores[0][i].item()
	label = g_labels[idx]
	if label not in seen_individuals:
	seen_individuals.add(label)
	unique_candidates.append((idx, score))

	# Fine search and visualization
	feats_q = extractor.extract(load_image(input_path).to(DEVICE))

	best_score = -1
	best_candidate_idx = -1

	# Initialize output lists (size 3)
	cand_logs = ["Waiting for data...", "Waiting for data...", "Waiting for data..."]
	cand_viz_paths = [None, None, None]

	for rank, (idx, arcface_sim) in enumerate(unique_candidates):
	path = g_paths[idx]
	label = g_labels[idx]
	species = g_species[idx]

	try:
	if not os.path.exists(path): continue

	feats_c = extractor.extract(load_image(path).to(DEVICE))
	with torch.no_grad():
	matches = matcher({"image0": feats_q, "image1": feats_c})
	matches = rbd(matches)

	geo_matches = len(matches["matches"])
	sim_percent = arcface_sim * 100

	# Create individual log string
	log_str = f"### Candidate {rank+1}: {species} / {label}\n"
	log_str += f"Coarse-Search Confidence: {sim_percent:.1f}%   \|   📐 Geometric Matches: {geo_matches}"
	cand_logs[rank] = log_str

	viz_name = f"viz_rank_{rank}"
	viz_path = create_match_visualization(
	input_path, path,
	feats_q['keypoints'], feats_c['keypoints'],
	matches['matches'], viz_name
	)
	cand_viz_paths[rank] = viz_path

	if geo_matches > best_score:
	best_score = geo_matches
	best_candidate_idx = idx

	except Exception as e:
	cand_logs[rank] = f"Error processing candidate: {e}"

	# Final decision calculation
	CONFIDENCE_THRESHOLD = 15
	if best_candidate_idx != -1 and best_score > CONFIDENCE_THRESHOLD:
	pred_species = g_species[best_candidate_idx]
	pred_id = g_labels[best_candidate_idx]
	is_correct = (pred_id == true_id)

	if true_id == "Unknown": header = f"### ❓ MATCH FOUND (No Ground Truth)\n"
	elif is_correct: header = f"### ✅ CORRECT MATCH!\n"
	else: header = f"### ❌ INCORRECT MATCH\n"

	header += f"Ground Truth: {true_species} / {true_id}   ➡️   Prediction: {pred_species} / {pred_id}\n"
	header += f"(Confirmed with {best_score} geometric keypoints)"
	else:
	header = "### ⚠️ UNKNOWN / NO MATCH\n"
	header += f"Ground Truth: {true_species} / {true_id}\n"
	header += f"Prediction: None (Best match only had {best_score} keypoints)\n"

	# Return: Header, then (Log, Img) for Cand 1, then (Log, Img) for Cand 2, etc.
	return (header,
	cand_logs[0], cand_viz_paths[0],
	cand_logs[1], cand_viz_paths[1],
	cand_logs[2], cand_viz_paths[2])

	# Setup for the user interface
	examples_list = []
	if os.path.exists(TEST_QUERIES_DIR):
	examples_list = [os.path.join(TEST_QUERIES_DIR, f) for f in os.listdir(TEST_QUERIES_DIR) if f.lower().endswith(('.jpg', '.png'))]

	with gr.Blocks(title="Wildlife Re-ID: Coarse-to-Fine Demo") as demo:
	gr.Markdown("# Wildlife Re-ID: Coarse-to-Fine Demo")
	gr.Markdown("Select a test image. The system finds the Top 3 UNIQUE individuals using embeddings, then verifies them using geometry.")

	with gr.Row():
	# Left Column: Input
	with gr.Column(scale=1):
	input_img = gr.Image(type="filepath", label="Test Image", height=300)
	gr.Examples(examples=examples_list, inputs=input_img, label="Test Examples", examples_per_page=4)
	submit_btn = gr.Button("Run Identification", variant="primary", size="lg")

	# Right Column: Vertical Stack of Candidates
	with gr.Column(scale=2):
	header_md = gr.Markdown(label="Final Decision")

	# Candidate 1 Group
	with gr.Group():
	log1 = gr.Markdown()
	# FIX: Removed show_download_button AND height
	img1 = gr.Image(label="Visualization", show_label=False)

	# Candidate 2 Group
	with gr.Group():
	log2 = gr.Markdown()
	img2 = gr.Image(label="Visualization", show_label=False)

	# Candidate 3 Group
	with gr.Group():
	log3 = gr.Markdown()
	img3 = gr.Image(label="Visualization", show_label=False)

	submit_btn.click(
	fn=predict,
	inputs=input_img,
	outputs=[header_md, log1, img1, log2, img2, log3, img3]
	)

	demo.launch(allowed_paths=[TEST_QUERIES_DIR])