Upload folder using huggingface_hub

app.py

@@ -0,0 +1,25 @@
+import gradio as gr
+import cv2
+from inference import InferencePipeline
+
+pipeline = InferencePipeline("artifacts/model.onnx")
+
+def predict(image):
+    if image.shape[0] > image.shape[1]:
+        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
+    image = cv2.resize(image, (640, 480), interpolation=cv2.INTER_AREA)
+    return pipeline(image)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Banknotes Classifier",)
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="numpy", label="Input Image")
+            predict_button = gr.Button("Predict")
+        with gr.Column():
+            output = gr.Textbox(label="Prediction")
+    predict_button.click(predict, inputs=image_input, outputs=output)
+
+if __name__ == "__main__":
+    demo.launch()

artifacts/.gitignore

@@ -0,0 +1 @@
+/model.onnx

artifacts/checkpoints/.gitignore

@@ -0,0 +1 @@
+/best-checkpoint.ckpt

artifacts/model.onnx

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

inference.py

@@ -0,0 +1,43 @@
+import numpy as np
+import cv2
+import onnxruntime as ort 
+
+
+class InferencePipeline:
+    def __init__(self, model_path):
+        self.model = ort.InferenceSession(model_path)
+        self.input_name = self.model.get_inputs()[0].name
+        self.output_name = self.model.get_outputs()[0].name
+        self.mapping = {
+            0: "1 pound face",
+            1: "1 pound back",
+            2: "5 pounds face",
+            3: "5 pounds back",
+            4: "10 pounds face",
+            5: "10 pounds back",
+            6: "20 pounds face",
+            7: "20 pounds back",
+            8: "50 pounds face",
+            9: "50 pounds back",
+            10: "100 pounds face",
+            11: "100 pounds back",
+            12: "200 pounds face",
+            13: "200 pounds back",
+        }
+
+    def __call__(self, image):
+        image = self._prepare_input(image)
+        output = self.predict(image)
+        predicted_class = self.mapping[output.argmax().item()]
+        predicted_confidence = output[0][output.argmax().item()] * 100
+        return f"The model is {predicted_confidence:.2f}% confident that the image is {predicted_class}"
+
+    def _prepare_input(self, image):
+        image = image.astype(np.float32)
+        image = image.transpose(2, 0, 1)
+        image = np.expand_dims(image, axis=0)
+        return image
+
+    def predict(self, image):
+        output = self.model.run([self.output_name], {self.input_name: image})[0]
+        return output

requirements.txt

@@ -0,0 +1,5 @@
+dvc==3.53.1
+dvc-gdrive
+numpy==1.26.0
+opencv_python_headless==4.10.0.84
+onnxruntime==1.18.1