script for eval

compute_seg_result.py

+from ultralytics import YOLO
+import cv2
+import matplotlib.pyplot as plt
+import argparse
+import os
+
+# Create the parser
+parser = argparse.ArgumentParser()
+# Add an argument
+parser.add_argument('--path_evaluation', type=str, required=True)
+parser.add_argument('--path_result', type=str, required=True)
+parser.add_argument('--class_object', type=str, required=True)
+parser.add_argument('--path_model_yolo', type=str, required=True)
+# Parse the argument
+args = parser.parse_args()
+
+
+# load the model.
+model = YOLO(f"{args.path_model_yolo}")
+#
+
+path_evaluation_data = args.path_evaluation + "/" + args.class_object + "/RGB_resized"
+
+for files in os.listdir(path_evaluation_data): 
+    #print("files : ", files)
+    #print(f"{path_evaluation_data}/{files}")
+
+    try:
+        results = model.predict(source=f"{path_evaluation_data}/{files}", conf=0.5, 
+save=True)
+    #print(results)
+        results1 = results[0].to('cpu')
+        results11 = results1.numpy()
+    # boxes = results11.boxes  # Boxes object for bbox outputs
+    # probs = results11.probs  # Class probabilities for classification outputs
+#
+        masks = results11.masks  # Masks object for segmentation masks outputs
+        #print(masks)
+        mask_res = masks.data[0]
+        plt.imsave(f"{args.path_result}/{files}", mask_res, cmap='gray') 
+        print("images saved : ", files )
+    except : 
+        print("no prediction")