import collections import heapq import json import os import logging import faiss import requests import gradio as gr import numpy as np import torch import torch.nn.functional as F from open_clip import create_model, get_tokenizer from torchvision import transforms from PIL import Image import io from pathlib import Path from huggingface_hub import hf_hub_download log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s" logging.basicConfig(level=logging.INFO, format=log_format) logger = logging.getLogger() hf_token = os.getenv("HF_TOKEN") model_str = "hf-hub:imageomics/bioclip" tokenizer_str = "ViT-B-16" txt_emb_npy = hf_hub_download(repo_id="pyesonekyaw/biome_lfs", filename='txt_emb_species.npy', repo_type="dataset") txt_names_json = "txt_emb_species.json" min_prob = 1e-9 k = 5 ranks = ("Kingdom", "Phylum", "Class", "Order", "Family", "Genus", "Species") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") preprocess_img = transforms.Compose( [ transforms.ToTensor(), transforms.Resize((224, 224), antialias=True), transforms.Normalize( mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711), ), ] ) MIN_PROB = 1e-9 TOP_K_PREDICTIONS = 5 TOP_K_CANDIDATES = 250 TOP_N_SIMILAR = 22 SIMILARITY_BOOST = 0.2 VOTE_THRESHOLD = 3 SIMILARITY_THRESHOLD = 0.99 # Add paths for RAG PHOTO_LOOKUP_PATH = f"./photo_lookup.json" SPECIES_LOOKUP_PATH = f"./species_lookup.json" theme = gr.themes.Base( primary_hue=gr.themes.colors.teal, secondary_hue=gr.themes.colors.blue, neutral_hue=gr.themes.colors.gray, text_size=gr.themes.sizes.text_lg, ).set( button_primary_background_fill="#114A56", button_primary_background_fill_hover="#114A56", block_title_text_weight="600", block_label_text_weight="600", block_label_text_size="*text_md", ) EXAMPLES_DIR = Path("examples") example_images = sorted(str(p) for p in EXAMPLES_DIR.glob("*.jpg")) def indexed(lst, indices): return [lst[i] for i in indices] def format_name(taxon, common): taxon = " ".join(taxon) if not common: return taxon return f"{taxon} ({common})" def combine_duplicate_predictions(predictions): """Combine predictions where one name is contained within another.""" combined = {} used = set() # Sort by length of name (longer names first) and probability items = sorted(predictions.items(), key=lambda x: (-len(x[0]), -x[1])) for name1, prob1 in items: if name1 in used: continue total_prob = prob1 used.add(name1) # Check remaining predictions for name2, prob2 in predictions.items(): if name2 in used: continue # Convert to lowercase for comparison name1_lower = name1.lower() name2_lower = name2.lower() # Check if one name contains the other if name1_lower in name2_lower or name2_lower in name1_lower: total_prob += prob2 used.add(name2) combined[name1] = total_prob # Normalize probabilities total = sum(combined.values()) return {k: v/total for k, v in combined.items()} @torch.no_grad() def open_domain_classification(img, rank: int, return_all=False): """ Predicts from the entire tree of life using RAG approach. """ logger.info(f"Starting open domain classification for rank: {rank}") img = preprocess_img(img).to(device) img_features = model.encode_image(img.unsqueeze(0)) img_features = F.normalize(img_features, dim=-1) # Get zero-shot predictions logits = (model.logit_scale.exp() * img_features @ txt_emb).squeeze() probs = F.softmax(logits, dim=0) # Get similar images votes and metadata species_votes, similar_images = get_similar_images_metadata(img_features, faiss_index, id_mapping, name_mapping) if rank + 1 == len(ranks): # Species level prediction topk = probs.topk(TOP_K_CANDIDATES) predictions = { format_name(*txt_names[i]): prob.item() for i, prob in zip(topk.indices, topk.values) } # Augment predictions with votes augmented_predictions = predictions.copy() for pred_name in predictions: pred_name_lower = pred_name.lower() for voted_species, vote_count in species_votes.items(): if voted_species in pred_name_lower or pred_name_lower in voted_species: augmented_predictions[pred_name] += SIMILARITY_BOOST * vote_count elif vote_count >= VOTE_THRESHOLD: augmented_predictions[voted_species] = vote_count * SIMILARITY_BOOST # Sort predictions sorted_predictions = dict(sorted( augmented_predictions.items(), key=lambda x: x[1], reverse=True )[:k]) # Normalize and combine duplicates total = sum(sorted_predictions.values()) sorted_predictions = {k: v/total for k, v in sorted_predictions.items()} sorted_predictions = combine_duplicate_predictions(sorted_predictions) logger.info(f"Top K predictions after combining duplicates: {sorted_predictions}") return sorted_predictions, similar_images # Higher rank prediction output = collections.defaultdict(float) for i in torch.nonzero(probs > MIN_PROB).squeeze(): output[" ".join(txt_names[i][0][: rank + 1])] += probs[i] # Incorporate votes for higher ranks for species, vote_count in species_votes.items(): try: # Find matching taxonomy in txt_names for taxonomy, _ in txt_names: if species in " ".join(taxonomy).lower(): higher_rank = " ".join(taxonomy[: rank + 1]) output[higher_rank] += SIMILARITY_BOOST * vote_count break except Exception as e: logger.error(f"Error processing vote for species {species}: {e}") # Get top-k predictions and normalize topk_names = heapq.nlargest(k, output, key=output.get) prediction_dict = {name: output[name] for name in topk_names} # Normalize probabilities to sum to 1 total = sum(prediction_dict.values()) prediction_dict = {k: v/total for k, v in prediction_dict.items()} prediction_dict = combine_duplicate_predictions(prediction_dict) logger.info(f"Prediction dictionary after combining duplicates: {prediction_dict}") return prediction_dict, similar_images def change_output(choice): return gr.Label(num_top_classes=k, label=ranks[choice], show_label=True, value=None) def get_cache_paths(name="demo"): """Get paths for cached FAISS index and ID mapping.""" return { 'index': hf_hub_download(repo_id="pyesonekyaw/biome_lfs", filename='cache/faiss_cache_demo.index', repo_type="dataset"), 'mapping': hf_hub_download(repo_id="pyesonekyaw/biome_lfs", filename='cache/faiss_cache_demo_mapping.json', repo_type="dataset") } def build_name_mapping(txt_names): """Build mapping between scientific names and common names.""" name_mapping = {} for taxonomy, common_name in txt_names: if not common_name: continue if len(taxonomy) >= 2: scientific_name = f"{taxonomy[-2]} {taxonomy[-1]}".lower() common_name = common_name.lower() name_mapping[scientific_name] = (scientific_name, common_name) name_mapping[common_name] = (scientific_name, common_name) return name_mapping def load_faiss_index(): """Load FAISS index from cache.""" cache_paths = get_cache_paths() logger.info("Loading FAISS index from cache...") index = faiss.read_index(cache_paths['index']) with open(cache_paths['mapping'], 'r') as f: id_mapping = json.load(f) return index, id_mapping def get_similar_images_metadata(img_embedding, faiss_index, id_mapping, name_mapping): """Get metadata for similar images using FAISS search.""" img_embedding_np = img_embedding.cpu().numpy() if img_embedding_np.ndim == 1: img_embedding_np = img_embedding_np.reshape(1, -1) # Search for more images than needed to account for filtered matches distances, indices = faiss_index.search(img_embedding_np, TOP_N_SIMILAR * 2) # Filter out near-exact matches valid_indices = [] valid_distances = [] valid_count = 0 for dist, idx in zip(distances[0], indices[0]): # For inner product similarity, the distance is already the similarity similarity = dist if similarity > SIMILARITY_THRESHOLD: continue valid_indices.append(idx) valid_distances.append(similarity) valid_count += 1 if valid_count >= TOP_N_SIMILAR: break species_votes = {} similar_images = [] for idx, similarity in zip(valid_indices[:5], valid_distances[:5]): # Only process top 5 for display similar_img_id = id_mapping[idx] try: species_names = id_to_species_info.get(similar_img_id) species_names = [name for name in species_names if name] processed_names = set() for species in species_names: if not species: continue name_tuple = name_mapping.get(species) if name_tuple: processed_names.add(name_tuple[0]) else: processed_names.add(species) for species in processed_names: species_votes[species] = species_votes.get(species, 0) + 1 # Store similar image info if the image file exists # if img_path and os.path.exists(img_path): similar_images.append({ 'id': similar_img_id, 'species': next(iter(processed_names)) if processed_names else 'Unknown', 'common_name': species_names[-1], 'similarity': similarity # Add similarity score }) except Exception as e: logger.error(f"Error processing JSON for image {similar_img_id}: {e}") continue return species_votes, similar_images if __name__ == "__main__": logger.info("Starting.") model = create_model(model_str, output_dict=True, require_pretrained=True) model = model.to(device) logger.info("Created model.") model = torch.compile(model) logger.info("Compiled model.") tokenizer = get_tokenizer(tokenizer_str) id_to_photo_url = json.load(open(PHOTO_LOOKUP_PATH)) id_to_species_info = json.load(open(SPECIES_LOOKUP_PATH)) logger.info(f"Loaded {len(id_to_photo_url)} photo mappings") logger.info(f"Loaded {len(id_to_species_info)} species mappings") # Load text embeddings and build name mapping txt_emb = torch.from_numpy(np.load(txt_emb_npy, mmap_mode="r")).to(device) with open(txt_names_json) as fd: txt_names = json.load(fd) # Build name mapping name_mapping = build_name_mapping(txt_names) # Build or load FAISS index with test IDs faiss_index, id_mapping = load_faiss_index() # Define process_output function before using it def process_output(img, rank): predictions, similar_imgs = open_domain_classification(img, rank) logger.info(f"Number of similar images found: {len(similar_imgs)}") images = [] labels = [] for img_info in similar_imgs: img_id = img_info['id'] img_url = id_to_photo_url.get(img_id) img_url = img_url.replace("square", "small") logger.info(f"Processing image URL: {img_url}") try: # Try fetching from URL first response = requests.get(img_url) if response.status_code == 200: try: img = Image.open(io.BytesIO(response.content)) images.append(img) except Exception as e: logger.info(f"Failed to load image from URL: {e}") images.append(None) else: logger.info(f"Failed to fetch image from URL: {response}") images.append(None) # Add label regardless of image load success label = f"**{img_info['species']}**" if img_info['common_name']: label += f" ({img_info['common_name']})" label += f"\nSimilarity: {img_info['similarity']:.3f}" label += f"\n[View on iNaturalist](https://www.inaturalist.org/observations/{img_id})" labels.append(label) except Exception as e: logger.error(f"Error processing image {img_id}: {e}") images.append(None) labels.append("") # Pad arrays if needed images += [None] * (5 - len(images)) labels += [""] * (5 - len(labels)) logger.info(f"Final number of images: {len(images)}") logger.info(f"Final number of labels: {len(labels)}") return [predictions] + images + labels with gr.Blocks(theme=theme) as app: # Add header with gr.Row(variant="panel"): with gr.Column(scale=1): gr.Image("image.jpg", elem_id="logo-img", show_label=False ) with gr.Column(scale=30): gr.Markdown("""Biome is a vision foundation model-powered tool customized to identify Singapore's local biodiversity.

**Developed by**: Pye Sone Kyaw - AI Engineer @ Multimodal AI Team - AI Practice - GovTech SG

Under the hood, Biome is using [BioCLIP](https://github.com/Imageomics/BioCLIP) augmented with multimodal search and retrieval to enhance its Singapore-specific biodiversity classification capabilities.

Biome work best when the organism is clearly visible and takes up a substantial part of the image. """) with gr.Row(variant="panel", elem_id="images_panel"): img_input = gr.Image( height=400, sources=["upload"], type="pil" ) with gr.Row(): with gr.Column(): with gr.Row(): gr.Examples( examples=example_images, inputs=img_input, label="Example Images" ) rank_dropdown = gr.Dropdown( label="Taxonomic Rank", info="Which taxonomic rank to predict. Fine-grained ranks (genus, species) are more challenging.", choices=ranks, value="Species", type="index", ) open_domain_btn = gr.Button("Submit", variant="primary") with gr.Column(): open_domain_output = gr.Label( num_top_classes=k, label="Prediction", show_label=True, value=None, ) # New section for similar images with gr.Row(variant="panel"): with gr.Column(): gr.Markdown("### Most Similar Images from Database") with gr.Row(): similar_images = [ gr.Image(label="Similar Image 1", height=200, show_label=True), gr.Image(label="Similar Image 2", height=200, show_label=True), gr.Image(label="Similar Image 3", height=200, show_label=True), gr.Image(label="Similar Image 4", height=200, show_label=True), gr.Image(label="Similar Image 5", height=200, show_label=True), ] with gr.Row(): similar_labels = [ gr.Markdown("Species 1"), gr.Markdown("Species 2"), gr.Markdown("Species 3"), gr.Markdown("Species 4"), gr.Markdown("Species 5"), ] rank_dropdown.change( fn=change_output, inputs=rank_dropdown, outputs=[open_domain_output] ) open_domain_btn.click( fn=process_output, inputs=[img_input, rank_dropdown], outputs=[open_domain_output] + similar_images + similar_labels, ) with gr.Row(variant="panel"): gr.Markdown(""" **Disclaimer**: This is a proof-of-concept demo for non-commercial purposes. No data is stored or used for any form of training, and all data used for retrieval are from [iNaturalist](https://inaturalist.org/). The adage of garbage in, garbage out applies here - uploading images not biodiversity-related will yield unpredictable results. """) app.queue(max_size=20) app.launch(share=False, enable_monitoring=False, allowed_paths=["/app/"])