Add yolov5/ to sys.path() for *.py subdir exec (ultralytics#2949)

* Add yolov5/ to sys.path() for *.py subdir exec * Update export.py
danny-schwartz7 · May 22, 2021 · 5949b6e · 5949b6e
1 parent 4890499
commit 5949b6e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/models/export.py b/models/export.py
@@ -7,8 +7,9 @@
 import argparse
 import sys
 import time
+from pathlib import Path
 
-sys.path.append('./')  # to run '$ python *.py' files in subdirectories
+sys.path.append(Path(__file__).parent.parent.absolute().__str__())  # to run '$ python *.py' files in subdirectories
 
 import torch
 import torch.nn as nn

diff --git a/models/yolo.py b/models/yolo.py
@@ -4,8 +4,9 @@
 import logging
 import sys
 from copy import deepcopy
+from pathlib import Path
 
-sys.path.append('./')  # to run '$ python *.py' files in subdirectories
+sys.path.append(Path(__file__).parent.parent.absolute().__str__())  # to run '$ python *.py' files in subdirectories
 logger = logging.getLogger(__name__)
 
 from models.common import *
@@ -267,7 +268,7 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
     # Create model
     model = Model(opt.cfg).to(device)
     model.train()
-                        
+
     # Profile
     # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 320, 320).to(device)
     # y = model(img, profile=True)