Added working model!

DummyModel.py

@@ -16,7 +16,7 @@ class DummyModel(torch.nn.Module):
         super().__init__()
 
     def forward(self, x:list):
-        return torch.softmax(torch.rand(len(x), 5), 1)
+        return torch.softmax(torch.rand(len(x), 5), 1), 0
 
     def __call__(self, x:list):
         return self.forward(x)

Model.py

@@ -1,13 +1,20 @@
+import os
 from nail_classification.inference import Inference
 
 
 class Model:
-    def __init__(self):
-        file_paths = [hf_hub_download("lfolle/DeepNAPSIModel", f"version_{v}.ckpt",
-                                use_auth_token=os.environ['DeepNAPSIModel']) for v in [10, 11, 12, 13, 14]]
+    def __init__(self, DEBUG):
+        if DEBUG:
+            base = r"C:\Users\follels\Documents\hand-ki-model-weights\DeepNAPSIModel\inference_checkpoints_v1"
+            file_paths = [os.path.join(base, f"version_{v}") for v in range(10, 15)]
+        else:
+            file_paths = [hf_hub_download("lfolle/DeepNAPSIModel", f"version_{v}.ckpt",
+                                    use_auth_token=os.environ['DeepNAPSIModel']) for v in [10, 11, 12, 13, 14]]
         self.inference = Inference(file_paths)
 
-
     def predict(self, x):
         y_hat, uncertainty = self.inference.predict(x)
         return y_hat, uncertainty
+
+    def __call__(self, x):
+        return self.predict(x)

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+from PIL import Image
 
 from backend import Infer
 
@@ -6,11 +7,16 @@ DEBUG = True
 
 infer = Infer(DEBUG)
 
-with gr.Blocks(analytics_enabled=False, title="DeepNAPSI Prediction") as demo:
+
+with gr.Blocks(analytics_enabled=False, title="DeepNAPSI") as demo:
     with gr.Column():
-        gr.Markdown("Upload an image of the hand and click **Predict NAPSI** to see the output.")
+        gr.Markdown("## Welcome to the DeepNAPSI application!")
+        gr.Markdown("Upload an image of the one hand and click **Predict NAPSI** to see the output.\n" \
+                    "> Note: Make sure there are no identifying information present in the image. The prediction can take up to 1 minute.")
         with gr.Column():
-            image_input = gr.Image()
+            with gr.Row():
+                image_input = gr.Image()
+                example_images = gr.Examples(["assets/example_1.jpg", "assets/example_2.jpg", "assets/example_3.jpg"], image_input)
             with gr.Row():
                 image_button = gr.Button("Predict NAPSI")
             outputs = []

backend.py

@@ -5,11 +5,13 @@ from huggingface_hub import hf_hub_download
 from nail_detection.main import get_nails
 
 from DummyModel import load_dummy_model
+from Model import Model
 
 
 class Infer():
     def __init__(self, DEBUG):
-        self.model = load_dummy_model(DEBUG)
+        # self.model = load_dummy_model(DEBUG)
+        self.model = Model(DEBUG)
 
     def predict(self, data):
         nails = get_nails(cv2.cvtColor(data, cv2.COLOR_RGB2BGR))
@@ -20,7 +22,9 @@ class Infer():
                 predictions.append(np.zeros((64, 64, 3)))
                 predictions.append(-1)
         else:
-            napsi_predictions = torch.argmax(self.model(nails), 1)
+            model_prediction, uncertainty = self.model(nails)
+            model_prediction = model_prediction[0]
+            napsi_predictions = torch.argmax(model_prediction, 1)
             napsi_sum = int(napsi_predictions.sum().detach().cpu())
             predictions.append(napsi_sum)
             for napsi_prediction, nail in zip(napsi_predictions, nails):

requirements.txt

@@ -1,6 +1,8 @@
-cvs2
+opencv-python
 numpy
-git+https://lfolle:${HANDKIMODEL}@huggingface.co/lfolle/handkimodel
+git+https://git5.cs.fau.de/folle/hand-ki-model.git
 gradio
 huggingface_hub
-torch
+--extra-index-url https://download.pytorch.org/whl/cu116
+torch
+torchvision