Create app.py

app.py

@@ -0,0 +1,69 @@
+import numpy as np
+import gradio as gr
+from PIL import Image
+import onnxruntime as ort
+
+def resize_and_crop(image, image_size):
+    
+    # Resize the image such that the shortest side is image_size
+    original_size = image.size
+    ratio = float(image_size) / min(original_size)
+    new_size = tuple([int(x * ratio) for x in original_size])
+    resized_image = image.resize(new_size, Image.LANCZOS)
+    
+    # Calculate coordinates for center cropping
+    left = (resized_image.width - image_size) / 2
+    top = (resized_image.height - image_size) / 2
+    right = (resized_image.width + image_size) / 2
+    bottom = (resized_image.height + image_size) / 2
+    
+    # Crop the image to image_size x image_size
+    cropped_image = resized_image.crop((left, top, right, bottom))
+
+    array = np.array(cropped_image)
+    array = np.transpose(array, (2, 0, 1))
+    array = np.expand_dims(array, axis=0)
+    array = (array/255).astype(np.float32)
+    
+    return array
+
+# Read class labels from a text file
+def read_labels(file_path):
+    with open(file_path, 'r') as f:
+        labels = [line.strip() for line in f.readlines()]
+    return labels
+
+# Load the class labels
+class_labels = read_labels('vocab_formatted.txt')
+
+taxon_included = read_labels('taxon_included.txt')
+string_taxon_included = ', '.join(sorted(taxon_included))
+
+taxon_not_included = read_labels('taxon_not_included.txt')
+string_taxon_not_included = ', '.join(sorted(taxon_not_included))
+
+# Load the ONNX model
+onnx_model = ort.InferenceSession('convnext_tiny.onnx')
+
+input_name = onnx_model.get_inputs()[0].name
+output_name = onnx_model.get_outputs()[0].name
+
+# Define the inference function
+def classify_image(image):
+    input_array = resize_and_crop(image, 320)
+    outputs = onnx_model.run([output_name], {input_name: input_array})[0]
+    result = {taxon: prob for taxon, prob in zip(class_labels, outputs[0])}
+    return result
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=5),
+    title="Image Classification for Freshwater Fish Species of Denmark",
+    description=f"**Upload an image to classify it**.\n\nSpecies included (Danish common name):\n*{string_taxon_included}*\n\nSpecies not included (yet!):\n*{string_taxon_not_included}*",
+)
+
+# Launch the app
+if __name__ == "__main__":
+    iface.launch()