Update test.py

test.py

@@ -21,26 +21,23 @@ print("Torchvision Version: ",torchvision.__version__)
 
 parser = argparse.ArgumentParser('arguments for testing the model')
 
-parser.add_argument('--ts_empty_folder', type=str, default="/data/taulukot/solukuvat/empty/test/",
+parser.add_argument('--ts_empty_folder', type=str, default="/path/to/empty/test/data/",
                     help='path to test data')
-parser.add_argument('--ts_ok_folder', type=str, default="/data/taulukot/solukuvat/ok/test/",
+parser.add_argument('--ts_ok_folder', type=str, default="/path/to/non-empty/test/data/",
                     help='path to test data')
-parser.add_argument('--results_folder', type=str, default="./results/aug_28022024/",
+parser.add_argument('--results_folder', type=str, default="./results/",
                     help='Folder for saving results')
-parser.add_argument('--model_path', type=str, default="/koodit/table_segmentation/empty_cell_detection/train/models/aug_b32_lr0001_28022024.onnx",
+parser.add_argument('--model_path', type=str, default="/path/to/model.onnx",
                     help='path to load model file from')
 parser.add_argument('--batch_size', type=int, default=16,
                     help='batch_size')
 parser.add_argument('--num_classes', type=int, default=2,
                     help='number of classes for classification')
-parser.add_argument('--name', type=str, default='empty_cell_augment_28022024',
+parser.add_argument('--name', type=str, default='test',
                     help='name given to result files')
 
 start = time.time()
 
-# nohup python test.py > logs/aug_test_28022024.txt 2>&1 &
-# echo $! > output/save_pid.txt
-
 torch.manual_seed(67)
 random.seed(67)
 
@@ -49,10 +46,8 @@ args = parser.parse_args()
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 Image.MAX_IMAGE_PIXELS = None
 
-# https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html
-
-
 def get_data():
+    """Combines test data paths and labels"""
     empty_path = Path(args.ts_empty_folder)
     ok_path = Path(args.ts_ok_folder)
 
@@ -62,11 +57,6 @@ def get_data():
     empty_labels = np.zeros(len(empty_files))
     ok_labels = np.ones(len(ok_files))
 
-    #ts_data_files = ts_data_files[:20]
-    #ts_data_labels = ts_data_labels[:20]
-    #ts_ok_files = ts_ok_files[:20]
-    #ts_ok_labels = ts_ok_labels[:20]
-
     ts_files = empty_files + ok_files
     ts_labels = np.concatenate((empty_labels, ok_labels))
 
@@ -77,12 +67,14 @@ def get_data():
 
 
 def initialize_model():
+    """Initializes .onnx model."""
     model = onnxruntime.InferenceSession(args.model_path)
     input_size = 224
     return model, input_size
 
-# Function for getting precision, recall and F-score metrics
+
 def get_precision_recall(y_true, y_pred):
+    """Calculates precision, recall and F-score metrics."""
     precision_recall_fscore = precision_recall_fscore_support(y_true, y_pred, average=None)
 
     prec_0 = precision_recall_fscore[0][0]
@@ -103,6 +95,7 @@ def get_precision_recall(y_true, y_pred):
 
 
 def createConfusionMatrix(y_true, y_pred):
+    """Creates confusion matrix based on the predicted and true labels."""
     classes = np.array(['empty', 'ok'])
 
     # Build confusion matrix
@@ -114,6 +107,7 @@ def createConfusionMatrix(y_true, y_pred):
     return sn.heatmap(df_cm, annot=True).get_figure()
 
 def save_preds(y_true, y_pred, paths):
+    """Saves file names and labels of incorrectly classified images."""
     # Identifies images that were not classified correctly
     incorrect_indices = np.where(y_true != y_pred)
     incorrectly_predicted_images = paths[incorrect_indices]
@@ -122,15 +116,12 @@ def save_preds(y_true, y_pred, paths):
 
     print(f'{len(incorrect_preds)} incorrect predictions')
 
-    # Save file names and labels of incorrectly classified images
     with open(args.results_folder + args.name + '_incorrect_preds', "w") as fp:
         json.dump(incorrect_preds, fp)
 
 # Initialize the model for this run
 model, input_size = initialize_model()
 
-# Print the model we just instantiated
-#print(model_ft)
 
 data_transforms = transforms.Compose([
         transforms.Resize((input_size, input_size)),
@@ -141,8 +132,8 @@ print("Initializing Datasets and Dataloaders...")
 
 ts_files, ts_labels = get_data()
 
-# Function for getting model predictions on test data
 def test_model(model, ts_files, ts_labels):
+    """Get model predictions on test data."""
     since = time.time()
     label_preds = []
     true_labels = []
@@ -175,14 +166,11 @@ ts_labels = np.array(ts_labels)
 
 # Test model
 y_pred, y_true, paths = test_model(model, ts_files, ts_labels)
-# Saves information of incorrect predictions
+# Save information of incorrect predictions
 save_preds(y_true, y_pred, paths)
-# Calculates and prints precision, recall and F-score metrics
+# Calculate and print precision, recall and F-score metrics
 get_precision_recall(y_true, y_pred)
 
-# Save confusion matrix to Tensorboard
-#cm = createConfusionMatrix(y_true, y_pred)
-#writer.add_figure("Confusion matrix", cm)
 # Create and save confusion matrix of the predictions and true labels
 conf_matrix = ConfusionMatrixDisplay.from_predictions(y_true, y_pred, normalize='true', display_labels=np.array(['empty', 'ok']))
 plt.savefig(args.results_folder + args.name + '_conf_matrix.jpg', bbox_inches='tight')