update model code

model.py

@@ -137,6 +137,11 @@ def check_file(image_path):
                 outputs = model(images)
                 _, predicted = torch.max(outputs.data, 1)
                 output = predicted
+                # Convert logits to probabilities
+                probabilities = F.softmax(outputs, dim=1)
+                # Get confidence score and prediction
+                confidence, d = torch.max(probabilities, 1)
+                print(confidence)
                 return predicted
 
     def remove_module_from_checkpoint(checkpoint):
@@ -157,4 +162,4 @@ def check_file(image_path):
     model = nn.DataParallel(model)
     output = test_model(model, test_loader, device)
     print(output)
-    return "Control" if output.item() == 0 else "Axial"
+    return "control" if output.item() == 0 else "ms"