updates

detect.py

@@ -8,10 +8,18 @@ from utils.utils import *
 from utils import torch_utils
 
 
-def detect(cfg, weights, images, output='output', img_size=416, conf_thres=0.3, nms_thres=0.45, save_txt=False,
-           save_images=True):
+def detect(
+        cfg,
+        weights,
+        images,
+        output='output',
+        img_size=416,
+        conf_thres=0.3,
+        nms_thres=0.45,
+        save_txt=False,
+        save_images=True
+):
     device = torch_utils.select_device()
-
     os.system('rm -rf ' + output)
     os.makedirs(output, exist_ok=True)
 
@@ -39,7 +47,6 @@ def detect(cfg, weights, images, output='output', img_size=416, conf_thres=0.3,
         t = time.time()
 
         # Get detections
-        with torch.no_grad():
         img = torch.from_numpy(img).unsqueeze(0).to(device)
         if ONNX_EXPORT:
             pred = torch.onnx._export(model, img, 'weights/model.onnx', verbose=True)
@@ -92,4 +99,12 @@ if __name__ == '__main__':
     opt = parser.parse_args()
     print(opt)
 
-    detect(opt.cfg, opt.weights, opt.images, img_size=opt.img_size, conf_thres=opt.conf_thres, nms_thres=opt.nms_thres)
+    with torch.no_grad():
+        detect(
+            opt.cfg,
+            opt.weights,
+            opt.images,
+            img_size=opt.img_size,
+            conf_thres=opt.conf_thres,
+            nms_thres=opt.nms_thres
+        )