Upload predict.py

handetect/predict.py

@@ -6,17 +6,13 @@ from PIL import Image
 from handetect.models import *
 from torchmetrics import ConfusionMatrix
 import matplotlib.pyplot as plt
-import pathlib
-import sys
 
 # Define the path to your model checkpoint
 model_checkpoint_path = "model.pth"
 
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-NUM_CLASSES = len(
-    os.listdir(r"C:\Users\User\Documents\PISTEK\HANDETECT\data\train\Task 1")
-)  # Update with the correct number of classes
+NUM_CLASSES = 6
 
 # Define transformation for preprocessing the input image
 preprocess = transforms.Compose(
@@ -41,7 +37,7 @@ torch.set_grad_enabled(False)
 
 def predict_image(image_path, model=model, transform=preprocess):
     # Define images variable to recursively list all the data file in the image_path
-    classes = os.listdir(r"C:\Users\User\Documents\PISTEK\HANDETECT\data\train\Task 1")
+    classes = ['Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease']
 
     print("---------------------------")
     print("Image path:", image_path)