add app

app.py

@@ -0,0 +1,142 @@
+import json
+import math
+import sqlite3
+
+import streamlit as st
+import torch
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from torchvision import transforms
+from transformers import AutoModelForImageClassification
+
+# Set the page title
+st.title("Global Bird Classification App")
+
+# Upload an image
+uploaded_file = st.file_uploader("Please select an image", type=["jpg", "jpeg", "png"])
+
+# Input latitude and longitude (optional)
+latitude = st.number_input("Enter latitude (optional)", value=None, format="%f")
+longitude = st.number_input("Enter longitude (optional)", value=None, format="%f")
+
+lang = st.selectbox(
+    "Result Language",
+    options=[2, 1, 0],
+    format_func=lambda x: {
+        2: "Latina (Nomen Scientificum)",
+        1: "English (IOC 10.1)",
+        0: "中文 （中国大陆）",
+    }[x]
+)
+
+
+classify_transforms = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+# crop and classification
+def classify_objects(classification_model, image, species_list):
+    input_tensor = classify_transforms(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        logits = classification_model(input_tensor)[0]
+        filtered = get_filtered_predictions(logits, species_list)
+        return softmax(filtered)
+
+
+def softmax(tuples):
+    # `torch.nn.functional.softmax` requires the input to be `Tensor`, so I implemented it myself
+    values = [t[1] for t in tuples]
+    exp_values = [math.exp(v) for v in values]
+    sum_exp_values = sum(exp_values)
+    softmax_values = [ev / sum_exp_values for ev in exp_values]
+    updated_tuples = [(t[0], softmax_values[i]) for i, t in enumerate(tuples)]
+    updated_tuples.sort(key=lambda t: t[1], reverse=True)
+
+    return updated_tuples
+
+
+def get_filtered_predictions(predictions: list[float], species_list: list[int]) -> list[tuple[int, float]]:
+    original = {index: value for index, value in enumerate(predictions)}
+
+    if species_list:
+        filtered_predictions = [(key, value) for key, value in original.items() if key in species_list]
+    else:
+        filtered_predictions = [(key, value) for key, value in original.items()]
+
+    return filtered_predictions
+
+
+class DistributionDB:
+    def __init__(self, db_path):
+        self.con = sqlite3.connect(db_path)
+        self.cur = self.con.cursor()
+
+
+    def get_list(self, lat, lng) -> list:
+        self.cur.execute(f'''
+SELECT m.cls
+FROM distributions AS d
+LEFT OUTER JOIN places AS p
+  ON p.worldid = d.worldid
+LEFT OUTER JOIN sp_cls_map AS m
+  ON d.species = m.species
+WHERE p.south <= {lat}
+  AND p.north >= {lat}
+  AND p.east >= {lng}
+  AND p.west <= {lng}
+GROUP BY d.species, m.cls;
+''')
+
+        return [row[0] for row in self.cur]
+
+    def close(self):
+        self.cur.close()
+        self.con.close()
+
+
+# If the user uploads an image
+if uploaded_file is not None:
+    try:
+        sqlite_path = hf_hub_download(repo_id='sunjiao/osea', filename='avonet.db')
+        st.success(f"Successfully downloaded distribution database from Hugging Face Hub!")
+
+        label_map_path = hf_hub_download(repo_id='sunjiao/osea', filename='bird_info.json')
+        st.success(f"Successfully downloaded labels from Hugging Face Hub!")
+    except Exception as e:
+        st.error(f"Failed to download the file: {e}")
+        st.stop()
+
+    db = DistributionDB(sqlite_path)
+    species_list = db.get_list(latitude, longitude)
+    db.close()
+
+    # Open the image
+    image = Image.open(uploaded_file)
+
+    # Display the uploaded image
+    st.image(image, caption="Uploaded Image", use_container_width=True)
+
+    model = AutoModelForImageClassification.from_pretrained('sunjiao/osea')
+
+    results = classify_objects(model, image, species_list)
+
+    top3_results = results[:3]
+
+    with open(label_map_path, 'r') as f:
+        data = f.read()
+
+    bird_info = json.loads(data)
+
+    # Display the top 3 results and their probabilities
+    st.subheader("Classification Results (Top 3):")
+    for result in top3_results:
+        st.write(f"{bird_info[result[0]][lang]}: {result[1]:.4f}")
+
+# Display latitude and longitude if provided
+if latitude is not None and longitude is not None:
+    st.write(f"Entered Latitude: {latitude}, Longitude: {longitude}")