Added napsi sum, small refactoring.

DummyModel.py

@@ -2,12 +2,21 @@ import torch
 import torch.nn
 
 
+def load_dummy_model(DEBUG):
+    model = DummyModel()
+    if not DEBUG:
+        file_path = hf_hub_download("lfolle/DeepNAPSIModel", "dummy_model.pth",
+                                    use_auth_token=os.environ['DeepNAPSIModel'])
+        model.load_state_dict(torch.load(file_path))
+    return model
+
+
 class DummyModel(torch.nn.Module):
     def __init__(self):
         super().__init__()
 
-    def forward(self, x):
-        return torch.softmax(torch.rand(5), 0)
+    def forward(self, x:list):
+        return torch.softmax(torch.rand(len(x), 5), 1)
 
-    def __call__(self, x):
-        return self.forward(x)
+    def __call__(self, x:list):
+        return self.forward(x)

Model.py

@@ -0,0 +1,13 @@
+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]]
+        self.inference = Inference(file_paths)
+
+
+    def predict(self, x):
+        y_hat, uncertainty = self.inference.predict(x)
+        return y_hat, uncertainty

app.py

@@ -15,6 +15,7 @@ with gr.Blocks(analytics_enabled=False, title="DeepNAPSI Prediction") as demo:
                 image_button = gr.Button("Predict NAPSI")
             outputs = []
             with gr.Row():
+                outputs.append(gr.Number(label="DeepNAPSI Sum"))
                 with gr.Column():
                     outputs.append(gr.Image())
                     outputs.append(gr.Number(label="DeepNAPSI Thumb"))

backend.py

@@ -4,31 +4,26 @@ import numpy as np
 from huggingface_hub import hf_hub_download
 from nail_detection.main import get_nails
 
-from DummyModel import DummyModel
-
-
-def load_model(DEBUG):
-    model = DummyModel()
-    if not DEBUG:
-        file_path = hf_hub_download("lfolle/DeepNAPSIModel", "dummy_model.pth",
-                                    use_auth_token=os.environ['DeepNAPSIModel'])
-        model.load_state_dict(torch.load(file_path))
-    return model
+from DummyModel import load_dummy_model
 
 
 class Infer():
     def __init__(self, DEBUG):
-        self.model = load_model(DEBUG)
+        self.model = load_dummy_model(DEBUG)
 
     def predict(self, data):
         nails = get_nails(cv2.cvtColor(data, cv2.COLOR_RGB2BGR))
         predictions = []
         if nails is None:
+            predictions.append(-1)
             for _ in range(5):
                 predictions.append(np.zeros((64, 64, 3)))
                 predictions.append(-1)
         else:
-            for nail in nails:
+            napsi_predictions = torch.argmax(self.model(nails), 1)
+            napsi_sum = int(napsi_predictions.sum().detach().cpu())
+            predictions.append(napsi_sum)
+            for napsi_prediction, nail in zip(napsi_predictions, nails):
                 predictions.append(nail)
-                predictions.append(int(torch.argmax(self.model(nail))))
+                predictions.append(napsi_prediction)
         return predictions