Update app.py

app.py

@@ -1,12 +1,21 @@
 import gradio as gr
 import tensorflow as tf
 from tensorflow.keras.models import load_model
-from tensorflow.keras.preprocessing.image import img_to_array, load_img
+from tensorflow.keras.preprocessing.image import img_to_array
 import numpy as np
 import os
 import requests
 
-# Load your model and tokenizer
+# Define model URLs
+model_urls = {
+    'class': 'https://huggingface.co/spaces/sowbaranika13/ohio_space/resolve/main/inceptionv3_class.h5',
+    'serpentes': 'https://huggingface.co/spaces/sowbaranika13/ohio_space/resolve/main/inceptionv3_serpentes.h5',
+    'mammalia': 'https://huggingface.co/spaces/sowbaranika13/ohio_space/resolve/main/inceptionv3_mammalia.h5',
+    'aves': 'https://huggingface.co/spaces/sowbaranika13/ohio_space/resolve/main/inceptionv3_aves.h5',
+    'amphibia': 'https://huggingface.co/spaces/sowbaranika13/ohio_space/resolve/main/inceptionv3_amphibia.h5'
+}
+
+# Labels for class and species
 labels = {
     'class': ['amphibia', 'aves', 'invertebrates', 'lacertilia', 'mammalia', 'serpentes', 'testudines'],
     'serpentes': ["Butler's Gartersnake", "Dekay's Brownsnake", 'Eastern Gartersnake', 'Eastern Hog-nosed snake', 'Eastern Massasauga', 'Eastern Milksnake', 'Eastern Racer Snake', 'Eastern Ribbonsnake', 'Gray Ratsnake', "Kirtland's Snake", 'Northern Watersnake', 'Plains Gartersnake', 'Red-bellied Snake', 'Smooth Greensnake'],
@@ -15,10 +24,26 @@ labels = {
     'amphibia': ['American Bullfrog', 'American Toad', 'Green Frog', 'Northern Leopard Frog']
 }
 
+# Download and save models locally
+def download_model(url, model_path):
+    if not os.path.exists(model_path):
+        response = requests.get(url)
+        with open(model_path, 'wb') as file:
+            file.write(response.content)
+
+# Ensure the models directory exists
+os.makedirs("models", exist_ok=True)
+
+# Download all models
+for label, url in model_urls.items():
+    model_path = os.path.join("models", f"inceptionv3_{label}.h5")
+    download_model(url, model_path)
+
+# Load models
 hierarchical_models = {}
 for label in labels:
-    model_url = f"https://huggingface.co/spaces/sowbaranika13/ohio_space/resolve/main/inceptionv3_{label}.h5"
-    hierarchical_models[label] = load_model(model_url)
+    model_path = os.path.join("models", f"inceptionv3_{label}.h5")
+    hierarchical_models[label] = load_model(model_path)
 
 def load_and_preprocess_image(image, target_size=(299, 299)):
     img_array = img_to_array(image)
@@ -37,11 +62,10 @@ def predict(image):
     results['class'] = class_label
     
     # Predict species level
-    if class_label in ['serpentes', 'mammalia', 'aves', 'amphibia']:
-        species_preds = hierarchical_models[class_label].predict(image_array)
-        species_idx = np.argmax(species_preds)
-        species_label = labels[class_label][species_idx]
-        results['species'] = species_label
+    species_preds = hierarchical_models[class_label].predict(image_array)
+    species_idx = np.argmax(species_preds)
+    species_label = labels[class_label][species_idx]
+    results['species'] = species_label
     
     return results