make code more flexible

neural_network/inference.py

@@ -10,6 +10,7 @@ parser.add_argument('-i', '--input',
     help='path to the input image')
 args = vars(parser.parse_args())
 
+def main(path):
     # the computation device
     device = ('cuda' if torch.cuda.is_available() else 'cpu')
     # list containing all the class labels
@@ -38,9 +39,9 @@ transform = transforms.Compose([
 
 
     # read and preprocess the image
-image = cv2.imread(args['input'])
+    image = cv2.imread(path)
     # get the ground truth class
-gt_class = args['input'].split('/')[-2]
+    gt_class = path.split('/')[-2]
     orig_image = image.copy()
     # convert to RGB format
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
@@ -51,20 +52,8 @@ with torch.no_grad():
         outputs = model(image.to(device))
     output_label = torch.topk(outputs, 1)
     pred_class = labels[int(output_label.indices)]
-cv2.putText(orig_image, 
-    f"GT: {gt_class}",
-    (10, 25),
-    cv2.FONT_HERSHEY_SIMPLEX, 
-    0.6, (0, 255, 0), 2, cv2.LINE_AA
-)
-cv2.putText(orig_image, 
-    f"Pred: {pred_class}",
-    (10, 55),
-    cv2.FONT_HERSHEY_SIMPLEX, 
-    0.6, (0, 0, 255), 2, cv2.LINE_AA
-)
-print(f"GT: {gt_class}, pred: {pred_class}")
-cv2.imshow('Result', orig_image)
-cv2.waitKey(0)
-cv2.imwrite(f"outputs/{gt_class}{args['input'].split('/')[-1].split('.')[0]}.png",
-    orig_image)
+    
+    return pred_class
+
+if __name__ == "__main__":
+    main(args['input'])