Add Gradio and Documentation

README.md

@@ -10,4 +10,19 @@ pinned: false
 license: gpl-3.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Pytorch Resnet34 Bird Classification
+
+The project is an implementation of the ResNet34 model as per the [Microsoft Research Paper](https://arxiv.org/abs/1512.03385). The model is build using PyTorch and is trained on the [Birds Classification Dataset](https://www.kaggle.com/datasets/gpiosenka/100-bird-species) from Kaggle. 
+
+## 🚀 Getting Started
+
+All the code for training the model and exporting to ONNX format is present in the [notebook](notebooks) folder or you can use this [Kaggle Notebook](https://www.kaggle.com/gauthamkrishnan119/pytorch-resnet34-birds-classification) for training the model. It took ~1.5 hours to train the model on the complete dataset using a P100 GPU. The [app.py](app.py) file contains the code for deploying the model using Gradio.
+
+## 🤗 Demo
+
+You can try out the model on [Hugging Face Spaces](https://huggingface.co/spaces/gauthamk/pytorch-resnet34-bird-classification)
+
+## 🖥️ Sample Interface
+
+![Sample Inference](samples/sample1.png)
+![Sample Inference](samples/sample2.png)

app.py

@@ -0,0 +1,13 @@
+import gradio as gr
+import os
+from functions import *
+
+examples_dir = 'examples'
+title = "Birds Classification - ResNet34 PyTorch"
+examples = [os.path.join(examples_dir, i) for i in os.listdir('examples')]
+
+interface = gr.Interface(fn=predict, inputs=gr.Image(type= 'numpy', shape=(224, 224)).style(height= 256),
+            outputs= gr.Label(num_top_classes= 5),
+            examples= examples, title= title, css= '.gr-box {background-color: rgb(230 230 230);}')
+
+interface.launch()

functions.py

@@ -0,0 +1,47 @@
+import numpy as np
+import json
+import onnxruntime as rt
+
+model_path = 'model/model.onnx'
+class_path = 'model/birds_name_mapping.json'
+
+normalise_means = [0.4914, 0.4822, 0.4465]
+normalise_stds = [0.2023, 0.1994, 0.2010]
+
+def normalise_image(image):
+    image = image.copy()
+    for i in range(3):
+        image[:, i, :, :] = (image[:, i, :, :] - normalise_means[i]) / normalise_stds[i]
+    return image
+
+def load_class_names():
+    with open(class_path, 'r') as f:
+        class_names = json.load(f)
+    return class_names
+
+def predict(inp_image):
+    
+    class_names = load_class_names()
+
+    image = inp_image
+    image = image.transpose((2, 0, 1))
+
+    image = image / 255.0
+    image = np.expand_dims(image, axis=0)
+    image = normalise_image(image)
+    image = image.astype(np.float32)
+
+    sess = rt.InferenceSession(model_path)
+
+    input_name = sess.get_inputs()[0].name
+    output_name = sess.get_outputs()[0].name
+
+    output = sess.run([output_name], {input_name: image})[0]
+    prob = np.exp(output) / np.sum(np.exp(output), axis=1, keepdims=True)
+
+    top5 = np.argsort(prob[0])[-5:][::-1]
+
+    class_probs = {class_names[str(i)]: float(prob[0][i]) for i in top5}
+    print(class_probs)
+
+    return class_probs