app.py

import os
import base64
import json
from typing import Optional, Dict, Any

import gradio as gr
import numpy as np
from open_clip import create_model_and_transforms, get_tokenizer
from PIL import Image
import requests
import torch

from logging_config import logger
from helpers import l2_normalize, encode_image

# Set your API Gateway URL below.
API_GATEWAY_URL = os.getenv(
    "API_GATEWAY_URL",
    ""
)

API_GATEWAY_API_KEY = os.getenv(
    "API_GATEWAY_API_KEY",
    ""
)

SPECIES_RANK_API_URL = os.getenv(
    "SPECIES_RANK_API_URL",
    ""
)

SPECIES_RANK_LICENSE_PARAMS = os.getenv(
    "SPECIES_RANK_LICENSE_PARAMS",
    ""
)

MODEL_NAME = os.getenv(
    "MODEL_NAME",
    "hf-hub:imageomics/bioclip"
)

# Load country code mappings from the JSON file
# Assumes the JSON file is located in the same directory as app.py.
logger.info("Loading country code mappings...")
country_codes_path = os.path.join(os.path.dirname(__file__), "country_codes.json")
with open(country_codes_path, "r") as f:
    country_code_mappings = json.load(f)
logger.info("Country code mappings loaded successfully.")

# Load BioCLIP Model from Hugging Face
logger.info("Loading model from Hugging Face...")
model, _, preprocess = create_model_and_transforms(MODEL_NAME)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
tokenizer = get_tokenizer(MODEL_NAME)
model = model.to(device)
logger.info(f"Model loaded on device successfully: {device}")

# Gradio App Function
def app_function(uploaded_image: Optional[np.ndarray], country: Optional[str]) -> Dict[str, Any]:
    """Main function for the Gradio app.
    
    Processes the uploaded image, performs semantic search, and returns a summary, species information, and HTML output.

    Args:
        uploaded_image (Optional[np.ndarray]): Uploaded image as a NumPy array.
        country (Optional[str]): Country for filtering the search results.
    
    Returns:
        Tuple[str, Optional[str], Optional[str], str]: Summary, proposed scientific name, proposed common name, and HTML output.
    """
    if uploaded_image is None:
        logger.error("app_function: No image uploaded.")
        return "No image uploaded", None, None, ""
    
    if country is None:
        logger.error("app_function: No country selected.")
        return "No country selected", None, None, ""
    
    try:
        image = Image.fromarray(uploaded_image)
    except Exception as e:
        logger.exception("app_function: Error processing image. Check if a valid image array is provided. Exception: %s", e)
        return f"Error processing image: {e}", None, None, ""
    
    try:
        query_embedding = np.array(encode_image(image=image, preprocess=preprocess, model=model, device=device))
        query_embedding = l2_normalize(query_embedding).tolist()
        logger.info("app_function: Image encoded successfully. Embedding length: %d", len(query_embedding))
    except Exception as e:
        logger.exception("app_function: Error encoding image. Uploaded image shape: %s. Exception: %s", getattr(uploaded_image, 'shape', 'N/A'), e)
        return f"Error encoding image: {e}", None, None, ""
    
    payload = {"query_embedding": query_embedding, "country_code": country_code_mappings.get(country, "")}
    headers = {"x-api-key": API_GATEWAY_API_KEY}
    logger.info("app_function: Calling API Gateway with payload (embedding sample: %s...)", query_embedding[:5])

    # Print the query embedding for debugging
    # print(query_embedding)

    try:
        response = requests.post(API_GATEWAY_URL, json=payload, headers=headers)
        logger.info("app_function: API Gateway responded with status code %d", response.status_code)
    except Exception as e:
        logger.exception("app_function: Exception during API Gateway call with payload: %s. Exception: %s", payload, e)
        return f"Error calling API: {e}", None, None, ""
    
    if response.status_code != 200:
        logger.error("app_function: API Gateway returned error %d - %s", response.status_code, response.text)
        return f"API error: {response.status_code} - {response.text}", None, None, ""
    
    try:
        body = response.json()
        logger.info("app_function: Successfully parsed API Gateway response as JSON.")

        # Print the response for debugging
        # print(response.text)
        # print(response.status_code)

        # If body is a string with a list, try to load it
        if isinstance(body, str):
            try:
                results = json.loads(body)
            except Exception:
                results = body
        else:
            results = body
    except Exception as e:
        logger.exception("app_function: Error decoding API Gateway response. Exception: %s", e)
        return f"Error decoding response: {e}", None, None, ""

    urls = []
    image_urls = []
    scientific_names = []
    common_names = []
    similarity_scores = []
    
    for res in results:
        urls.append(res.get("url", ""))
        image_urls.append(res.get("image_url", ""))
        scientific_names.append(res.get("scientific_name", "N/A"))
        common_names.append(res.get("common_name", "N/A"))
        similarity_scores.append(res.get("similarity", 0))
    
    proposed_scientific = scientific_names[0]
    proposed_common = common_names[0]
    summary = "Found top 5 similar wildlife images."
    
    logger.info("app_function: Performing taxonomic lookup for '%s'", proposed_scientific)
    try:
        response = requests.get(f"{SPECIES_RANK_API_URL}{proposed_scientific}{SPECIES_RANK_LICENSE_PARAMS}", timeout=10)
        response.raise_for_status()
        results = response.json().get("results", [])
        if isinstance(results, list) and results:
            taxonomic_data = results[0]
            if not isinstance(taxonomic_data, dict):
                taxonomic_data = {}
        else:
            taxonomic_data = {}

    except requests.exceptions.RequestException as e:
        logger.error("app_function: Taxonomic lookup failed for '%s'. Exception: %s", proposed_scientific, e)
        taxonomic_data = {}

    proposed_kingdom = taxonomic_data.get("kingdom", "N/A")
    proposed_phylum = taxonomic_data.get("phylum", "N/A")
    proposed_class = taxonomic_data.get("class", "N/A")
    proposed_order = taxonomic_data.get("order", "N/A")
    proposed_family = taxonomic_data.get("family", "N/A")
    proposed_genus = taxonomic_data.get("genus", "N/A")
    logger.info("app_function: Taxonomic lookup complete. HTTP status code: %d", response.status_code)

    # Build HTML output for the 5 boxes in horizontal arrangement.
    boxes_html = "<div style='display: flex; justify-content: space-around; flex-wrap: nowrap;'>"
    for url, image_url, sci, com, similarity_score in zip(urls, image_urls, scientific_names, common_names, similarity_scores):
        try:
            r = requests.get(image_url, timeout=5)
            if r.status_code == 200:
                encoded_img = base64.b64encode(r.content).decode("utf-8")
                # Wrap the image in a container to keep it within fixed dimensions.
                img_tag = f"""
                    <div style="width:200px; height:150px; overflow:hidden; display:flex; align-items:center; justify-content:center;">
                        <img src='data:image/jpeg;base64,{encoded_img}' style='max-width:100%; max-height:100%; object-fit: contain;'/>
                    </div>
                    """
            else:
                img_tag = """
                    <div style="width:200px; height:150px; background:#eee; display:flex; align-items:center; justify-content:center;">
                        Error loading image
                    </div>
                    """
        except Exception as e:
            logger.exception("app_function: Error loading image from URL: %s. Exception: %s", image_url, e)
            img_tag = """
                    <div style="width:200px; height:150px; background:#eee; display:flex; align-items:center; justify-content:center;">
                        Error loading image
                    </div>
                    """
        
        box = f"""
            <div style='text-align: center; margin: 10px; flex: 1; border: 1px solid #ccc; min-height: 250px; display: flex; flex-direction: column; align-items: center; justify-content: center;'>
                {img_tag}
                <div style='font-size: 12px; margin-top: 5px;'>
                    <div><a href="{url}" target="_blank">View on iNaturalist</a></div>
                    <div>Scientific: {sci}</div>
                    <div>Common: {com}</div>
                    <div>Similarity: {similarity_score:.2f}</div>
                </div>
            </div>
        """
        boxes_html += box
    boxes_html += "</div>"    

    logger.info("app_function: Results processed and returned to Gradio interface successfully.")

    return {
        "summary": summary,
        "scientific_name": proposed_scientific,
        "common_name": proposed_common,
        "boxes_html": boxes_html,
        "taxonomy": {
            "kingdom": proposed_kingdom,
            "phylum": proposed_phylum,
            "class": proposed_class,
            "order": proposed_order,
            "family": proposed_family,
            "genus": proposed_genus,
            }
    }

# Gradio Interface Using Blocks Layout
with gr.Blocks(title="Wildlife Semantic Search with BioCLIP") as demo:
    # Custom CSS to fix the display size of the uploaded image.
    gr.HTML(
        """
        <style>
            /* Force the uploaded image to fit within 300x300px while preserving aspect ratio */
            #fixedImage img {
                object-fit: contain;
                width: 300px;
                height: 300px;
            }
            /* Style the logo to remove whitespace */
            .logo-image {
                object-fit: cover;
                object-position: center;
                width: 100%;
                height: 100%;
                display: block;
                margin: 0;
                padding: 0;
            }
            /* Custom style for the submit button */
            .submit-button {
                background: linear-gradient(90deg, green 0%, green 70%, orange 100%) !important;
                color: white !important;
                font-weight: bold !important;
            }
        </style>
        """
    )
    
    # Row 1: Logo and Description in two columns.
    with gr.Row(variant="panel"):
        with gr.Column(scale=1):
            gr.Image("logo/logo.jpg", elem_classes=["logo-image"], show_label=False)
        with gr.Column(scale=30):
            gr.Markdown(
                """
                ### Welcome to Ecologist – an AI-powered biodiversity explorer!

                **Ecologist** identifies wildlife species found in your selected country from an uploaded photo. 
                
                Powered by multimodal image retrieval and visual encoding with [BioCLIP](https://huggingface.co/imageomics/bioclip), the system extracts features from the image and matches them against a specialized database of the country's diverse flora and fauna.
                
                Both scientific and common names are provided within seconds, along with visually similar images that offer context about the country's rich natural heritage.
                
                Ecologist is a step towards celebrating and preserving the world’s unique wildlife through AI.
                """
            )
            country_dropdown = gr.Dropdown(
                label="Select Country",
                choices=["Indonesia", "Malaysia", "Singapore", "Thailand"],
                value="Singapore"
            )

    # Row 2: Image Upload with a fixed display container.
    with gr.Row(variant="panel"):
        with gr.Column():
            image_input = gr.Image(type="numpy", label="Upload Wildlife Image", elem_id="fixedImage")
    
    # Row 3: Submit Button.
    submit_button = gr.Button("Submit", elem_classes=["submit-button"])

    with gr.Row(variant="panel"):
        with gr.Column():
            gr.Examples(
                examples=[
                    ["examples/boar.jpg"],
                    ["examples/crow.jpg"],
                    ["examples/dragonfly.jpg"],
                    ["examples/macque.jpg"],
                    ["examples/otter.jpg"],
                    ["examples/parrot.jpg"],
                    ["examples/squirrel.jpg"],
                ],
                inputs=image_input,
                outputs=None,
                label="Example Wildlife Images",
            )

    # Row 4: Proposed Species Output.
    with gr.Row(variant="panel"):
        with gr.Column():
            gr.Markdown("## Identified Species")
    
    with gr.Row(variant="panel"):
        with gr.Column():
            proposed_kingdom_output = gr.Textbox(label="1. Kingdom", placeholder="N/A")
        with gr.Column():
            proposed_phylum_output = gr.Textbox(label="2. Phylum", placeholder="N/A")
        with gr.Column():
            proposed_class_output = gr.Textbox(label="3. Class", placeholder="N/A")

    with gr.Row(variant="panel"):
        with gr.Column():
            proposed_order_output = gr.Textbox(label="4. Order", placeholder="N/A")
        with gr.Column():
            proposed_family_output = gr.Textbox(label="5. Family", placeholder="N/A")
        with gr.Column():
            proposed_genus_output = gr.Textbox(label="6. Genus", placeholder="N/A")

    with gr.Row(variant="panel"):
        with gr.Column():
            proposed_scientific_output = gr.Textbox(label="7. Species (Scientific Name)", placeholder="No name yet")
        with gr.Column():
            proposed_common_output = gr.Textbox(label="8. Common Name", placeholder="No name yet")

    # Row 5: Pre-populated placeholder for 5 columns with borders.
    with gr.Row(variant="panel"):
        with gr.Column():
            gr.Markdown("## Most Similar Wildlife Images from Database")

    placeholder_boxes = "<div style='display: flex; justify-content: space-around; flex-wrap: nowrap;'>"
    for _ in range(5):
        placeholder_boxes += """
            <div style='text-align: center; margin: 10px; flex: 1; border: 1px solid #ccc; min-height: 250px; display: flex; align-items: center; justify-content: center;'>
                No image yet
            </div>
        """
    placeholder_boxes += "</div>"

    with gr.Row(variant="panel"):
        with gr.Column():
            html_output = gr.HTML(value=placeholder_boxes, container=True)

    with gr.Row(variant="panel"):
        with gr.Column():
            gr.Markdown(
                """
                **Disclaimer:**  
                Intended for non-commercial use, no user data is stored or used for training purposes, and all retrieval data is sourced from [iNaturalist](https://inaturalist.org/) and the [Global Biodiversity Information Facility (GBIF)](https://techdocs.gbif.org/en/openapi/). Results may vary depending on the input image. 
                
                **References:**  
                This project is inspired by the work on [Biome](https://huggingface.co/spaces/govtech/Biome) from GovTech Singapore.  
                
                **Acknowledgments:**  
                Gratitude to [Dylan Chan](https://www.pexels.com/@dylan-chan-2880813/), [Jesper](https://www.pexels.com/@jesper-425001880/), [Mark Baldovino](https://www.pexels.com/@odlab2/), [Sane Noor](https://www.pexels.com/@norsan/), [Soumen Chakraborty](https://www.pexels.com/@soumen-chakraborty-363019169/), [Tony Wu](https://www.pexels.com/@tonywuphotography/) and [Zett Foto](https://www.pexels.com/@zett-foto-194587/) for their wildlife images in [Pexels](https://www.pexels.com/).
                """
                )
    
    # Wrapping the function to only forward the necessary outputs.
    def wrapper(uploaded_image, country):
        result  = app_function(uploaded_image, country)
        taxonomy = result.get("taxonomy", {})

        # Print the summary for debugging
        # print(result.get("summary", "Failed to get top 5 similar wildlife images."))

        return (
            result.get("scientific_name"),
            result.get("common_name"),
            result.get("boxes_html"),
            taxonomy.get("kingdom"),
            taxonomy.get("phylum"),
            taxonomy.get("class"),
            taxonomy.get("order"),
            taxonomy.get("family"),
            taxonomy.get("genus")
        )

    submit_button.click(
        fn=wrapper, 
        inputs=[
            image_input, 
            country_dropdown
        ], 
        outputs=[
            proposed_scientific_output, 
            proposed_common_output, 
            html_output, 
            proposed_kingdom_output, 
            proposed_phylum_output, 
            proposed_class_output, 
            proposed_order_output, 
            proposed_family_output, 
            proposed_genus_output
        ]
    )

if __name__ == "__main__":
    demo.launch()