Update main.py

main.py

@@ -11,30 +11,22 @@ def run_transfer_learning(dataset, epochs, batch_size, imgsz, patience, cache, p
         "--imgsz", str(imgsz),
         "--patience", str(patience),
         "--cache", cache,
+        "--pretrained", str(pretrained),
+        "--cos_lr", str(cos_lr),
+        "--profile", str(profile),
+        "--plots", str(plots),
+        "--resume", str(resume),
+        "--augment", str(augment),
         "--model", model,
         "--run", run
     ]
 
-    # Append boolean flags conditionally
-    if pretrained:
-        command.append("--pretrained")
-    if cos_lr:
-        command.append("--cos_lr")
-    if profile:
-        command.append("--profile")
-    if plots:
-        command.append("--plots")
-    if resume:
-        command.append("--resume")
-    if augment:
-        command.append("--augment")
-    
     # Use subprocess to run the script with the arguments
     subprocess.run(command, check=True)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run transfer learning with YOLO model.")
-    
+
     # Define all arguments with default values from your script
     parser.add_argument('--dataset', type=str, choices=["Birds-Nest", "Common-VALID", "Electric-Substation", "InsPLAD-det"], help='Dataset name to be used')
     parser.add_argument("--epochs", type=int, default=1000, help="Number of epochs")
@@ -56,17 +48,17 @@ if __name__ == "__main__":
     # Call the function to run transfer learning
     run_transfer_learning(
         dataset=args.dataset,
-        epochs=args.epochs, 
-        batch_size=args.batch, 
-        imgsz=args.imgsz, 
-        patience=args.patience, 
-        cache=args.cache, 
-        pretrained=args.pretrained, 
-        cos_lr=args.cos_lr, 
-        profile=args.profile, 
+        epochs=args.epochs,
+        batch_size=args.batch,
+        imgsz=args.imgsz,
+        patience=args.patience,
+        cache=args.cache,
+        pretrained=args.pretrained,
+        cos_lr=args.cos_lr,
+        profile=args.profile,
         plots=args.plots,
         resume=args.resume,
         augment=args.augment,
-        model=args.model, 
+        model=args.model,
         run=args.run
     )