Add files

.gitignore

@@ -0,0 +1,3 @@
+*.pth
+flagged
+__pycache__

README.md

@@ -7,6 +7,4 @@ sdk: gradio
 sdk_version: 3.23.0
 app_file: app.py
 pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---

app.py

@@ -0,0 +1,18 @@
+import gradio as gr
+import os
+from functions import *
+
+examples_dir = 'examples'
+title = "CvT: Substructure Prediction"
+
+examples = []
+for root, dirs, files in os.walk(examples_dir):
+    for file in files:
+        if file.endswith(".jpg"):
+            examples.append([os.path.join(root, file), os.path.join(root, file)])
+
+interface = gr.Interface(fn=predict, inputs=gr.Image(type= 'pil', shape=(256, 256), image_mode= 'L'),
+            outputs= gr.Label(num_top_classes=2), cache_examples= False,
+            examples= examples, title= title, css= '.gr-box {background-color: rgb(230 230 230);}')
+
+interface.launch()

functions.py

@@ -0,0 +1,20 @@
+import numpy as np
+import cv2
+import onnxruntime as rt
+from PIL import Image
+
+onnx_path = 'model/model.onnx'
+
+def predict(img):
+    session = rt.InferenceSession(onnx_path)
+    input_name = session.get_inputs()[0].name
+    output_name = session.get_outputs()[0].name
+    
+    img = np.array(img).astype(np.float32)
+    img = img.reshape(1, 1, 256, 256)
+    img = img / 255.0
+    pred = session.run([output_name], {input_name: img})[0]
+    pred = np.exp(pred) / np.sum(np.exp(pred), axis=1, keepdims=True)
+
+    class_probs = {'No Substructure': str(pred[0][0]), 'Substructure': str(pred[0][1])}
+    return class_probs

model/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf0384da75ea8db84808deef7fc139c1648569baa7db78105f870b63498692ea
+size 93167130