Update evaluate code

estimate_model.py

@@ -26,7 +26,7 @@ def predictor(model, img, mask, device):
     return prediction
 
 
-def run_pred(args, weights_path, img_path, roi_mask_path):
+def run_pred(args, model, weights_path, img_path, roi_mask_path):
     assert os.path.exists(weights_path), f"weights {weights_path} not found."
     assert os.path.exists(img_path), f"image {img_path} not found."
     assert os.path.exists(roi_mask_path), f"image {roi_mask_path} not found."
@@ -38,9 +38,6 @@ def run_pred(args, weights_path, img_path, roi_mask_path):
     device = args.device
     print("using {} device.".format(device))
 
-    # create model
-    model = UKAN_large(num_classes=args.nb_classes)
-
     # load weights
     model.load_state_dict(torch.load(weights_path, map_location='cpu')['model_state'])
     model.to(device)

train_gpu.py

@@ -345,11 +345,16 @@ def main(args):
             print('*********No improving mIOU, No saving checkpoint*********')
 
     if args.predict and utils.is_main_process():
+        model_pred = create_model(
+            args.model,
+            num_classes=args.nb_classes,
+            args=args
+        )
         print('*******************STARTING PREDICT*******************')
         weights_path = f'./{args.save_weights_dir}/{args.model}_best_model.pth'
         img_path = "/mnt/d/MedicalSeg/CVC-ClinicDB/Original/1.png"
         roi_mask_path = "/mnt/d/MedicalSeg/CVC-ClinicDB/Ground Truth/1.png"
-        run_pred(args, weights_path, img_path, roi_mask_path)
+        run_pred(args, model_pred, weights_path, img_path, roi_mask_path)