Update app.py

app.py

@@ -29,6 +29,14 @@ MODEL_NAME = os.getenv(
     "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)
@@ -38,13 +46,14 @@ model = model.to(device)
 logger.info(f"Model loaded on device successfully: {device}")
 
 # Gradio App Function
-def app_function(uploaded_image: Optional[np.ndarray]) -> Tuple[str, Optional[str], Optional[str], str]:
+def app_function(uploaded_image: Optional[np.ndarray], country: Optional[str]) -> Tuple[str, Optional[str], Optional[str], str]:
     """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.
@@ -53,6 +62,10 @@ def app_function(uploaded_image: Optional[np.ndarray]) -> Tuple[str, Optional[st
         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:
@@ -67,7 +80,7 @@ def app_function(uploaded_image: Optional[np.ndarray]) -> Tuple[str, Optional[st
         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": "sg"}
+    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])
 
@@ -205,17 +218,22 @@ with gr.Blocks(title="Wildlife Semantic Search with BioCLIP") as demo:
         with gr.Column(scale=30):
             gr.Markdown(
                 """
-                ### Welcome to Ecologist – Singapore's AI-powered biodiversity explorer!
+                ### Welcome to Ecologist – an AI-powered biodiversity explorer!
 
-                **Ecologist** identifies wildlife species found in Singapore from an uploaded photo. 
+                **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 Singapore's diverse flora and fauna.
+                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 Singapore's rich natural heritage.
+                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 island country’s unique wildlife through AI.
+                Ecologist is a step towards celebrating and preserving the world’s unique wildlife through AI.
                 """
             )
+            country_dropdown = gr.Dropdown(
+                label="Select Country",
+                choices=["Singapore"],
+                value="Singapore"
+            )
 
     # Row 2: Image Upload with a fixed display container.
     with gr.Row(variant="panel"):
@@ -276,7 +294,7 @@ with gr.Blocks(title="Wildlife Semantic Search with BioCLIP") as demo:
             gr.Markdown(
                 """
                 **Disclaimer:**  
-                Not intended for commercial use, no user data is stored or used for training purposes, and all retrieval data is sourced from [iNaturalist](https://inaturalist.org/). Results may vary depending on the input image. 
+                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/). 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.  
@@ -287,15 +305,15 @@ with gr.Blocks(title="Wildlife Semantic Search with BioCLIP") as demo:
                 )
     
     # Wrapping the function to only forward the necessary outputs.
-    def wrapper(uploaded_image):
-        summary, proposed_scientific, proposed_common, boxes_html = app_function(uploaded_image)
+    def wrapper(uploaded_image, country):
+        summary, proposed_scientific, proposed_common, boxes_html = app_function(uploaded_image, country)
 
         # Print the summary for debugging
         # print(summary)
 
         return proposed_scientific, proposed_common, boxes_html
 
-    submit_button.click(fn=wrapper, inputs=image_input, outputs=[proposed_scientific_output, proposed_common_output, html_output])
+    submit_button.click(fn=wrapper, inputs=[image_input, country_dropdown], outputs=[proposed_scientific_output, proposed_common_output, html_output])
 
 if __name__ == "__main__":
     demo.launch()