Update app.py

app.py

@@ -1,34 +1,92 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForImageClassification, AutoTokenizer
-from PIL import Image
+import tensorflow as tf
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.image import img_to_array, load_img
+import numpy as np
+import os
+import shutil
+from docx import Document
+from docx.shared import Inches
+import matplotlib.pyplot as plt
 
 # Load your model and tokenizer
-model_name = "inceptionv3_class.h5"
-model = AutoModelForImageClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+class_thresholds = {
+    'class': 0.6,
+    'amphibia': 0.1,
+    'aves': 0.5,
+    'mammalia': 0.8,
+    'serpentes': 0.9
+}
+
+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'],
+    'mammalia': ['American Mink', 'Brown Rat', 'Eastern Chipmunk', 'Eastern Cottontail', 'Long-tailed Weasel', 'Masked Shrew', 'Meadow Jumping Mouse', 'Meadow Vole', 'N. Short-tailed Shrew', 'Raccoon', 'Star-nosed mole', 'Striped Skunk', 'Virginia Opossum', 'White-footed Mouse', 'Woodchuck', 'Woodland Jumping Mouse'],
+    'aves': ['Common Yellowthroat', 'Gray Catbird', 'Indigo Bunting', 'Northern House Wren', 'Song Sparrow', 'Sora'],
+    'amphibia': ['American Bullfrog', 'American Toad', 'Green Frog', 'Northern Leopard Frog']
+}
+
+hierarchical_models = {}
+for label in labels:
+    model_path = f"inceptionv3_{label}.h5"
+    if os.path.exists(model_path):
+        hierarchical_models[label] = load_model(model_path)
+
+def load_and_preprocess_image(image_path, target_size=(299, 299)):
+    img = load_img(image_path, target_size=target_size)
+    img_array = img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0)
+    img_array = tf.keras.applications.inception_v3.preprocess_input(img_array)
+    return img_array
+
+def copy_images(source_dir, image_path, label, confidence, threshold):
+    if confidence < threshold:
+        human_review_dir = os.path.join(source_dir, 'human', label)
+        os.makedirs(human_review_dir, exist_ok=True)
+        human_review_path = os.path.join(human_review_dir, os.path.basename(image_path))
+        shutil.copy(image_path, human_review_path)
+    else:
+        output_path = os.path.join(source_dir, label)
+        os.makedirs(output_path, exist_ok=True)
+        output_dir = os.path.join(output_path, os.path.basename(image_path))
+        shutil.copy(image_path, output_dir)
+
+def process_images(input_d, level, dir_name):
+    input_dir = os.path.join(input_d, dir_name)
+    image_paths = [os.path.join(input_dir, fname) for fname in os.listdir(input_dir)]
+    images = np.vstack([load_and_preprocess_image(image_path) for image_path in image_paths])
+    predictions = hierarchical_models[level].predict(images)
+    for image_path, prediction in zip(image_paths, predictions):
+        predicted_class_index = np.argmax(prediction)
+        predicted_class_label = labels[level][predicted_class_index]
+        confidence = prediction[predicted_class_index]
+        copy_images(input_d, image_path, predicted_class_label, confidence, class_thresholds[level])
+    if level == "class":
+        for label in ['serpentes', 'mammalia', 'aves', 'amphibia']:
+            class_dir = os.path.join(input_d, label)
+            if os.path.exists(class_dir):
+                process_images(input_d, label, label)
 
 def predict(image):
-    # Preprocess the image
-    inputs = tokenizer(images=image, return_tensors="pt")
-    
-    # Predict
-    outputs = model(**inputs)
-    logits = outputs.logits
-    probs = torch.nn.functional.softmax(logits, dim=-1)
-    
-    # Get top 5 predictions
-    top_probs, top_labels = torch.topk(probs, 5)
-    
-    # Map labels to human-readable class names
-    class_names = [model.config.id2label[label.item()] for label in top_labels[0]]
-    return {name: prob.item() for name, prob in zip(class_names, top_probs[0])}
+    image_path = "temp_image.jpg"
+    image.save(image_path)
+    level = "class"
+    dir_name = "temp_images"
+    os.makedirs(dir_name, exist_ok=True)
+    shutil.copy(image_path, os.path.join(dir_name, os.path.basename(image_path)))
+    process_images(".", level, dir_name)
+    predictions = {}
+    for label in ['serpentes', 'mammalia', 'aves', 'amphibia']:
+        class_dir = os.path.join(".", label)
+        if os.path.exists(class_dir):
+            predictions[label] = os.listdir(class_dir)
+    return predictions
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=predict,
     inputs=gr.inputs.Image(type="pil"),
-    outputs=gr.outputs.Label(num_top_classes=5),
+    outputs=gr.outputs.Label(num_top_classes=7),
     title="Image Classification",
     description="Upload an image to classify it into species and class level."
 )