From be8603b2dd955ebbdd6328207436ed3acac34293 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 4 Dec 2018 19:17:03 +0100
Subject: [PATCH] updates

---
 detect.py    | 6 +++---
 test.py      | 6 +++---
 utils/gcp.sh | 3 +++
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/detect.py b/detect.py
index 5c1d9fb8..7953f30e 100755
--- a/detect.py
+++ b/detect.py
@@ -35,11 +35,11 @@ def main(opt):
     model = Darknet(opt.cfg, opt.img_size)
 
     weights_path = f_path + 'weights/yolov3.pt'
-    if weights_path.endswith('.weights'):  # saved in darknet format
-        load_weights(model, weights_path)
-    else:  # endswith('.pt'), saved in pytorch format
+    if weights_path.endswith('.pt'):  # pytorch format
         if weights_path.endswith('weights/yolov3.pt') and not os.path.isfile(weights_path):
             os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights_path)
+    else:  # darknet format
+        load_weights(model, weights_path)
 
         checkpoint = torch.load(weights_path, map_location='cpu')
         model.load_state_dict(checkpoint['model'])
diff --git a/test.py b/test.py
index 4d739993..59e725ca 100644
--- a/test.py
+++ b/test.py
@@ -35,12 +35,12 @@ def main(opt):
     model = Darknet(opt.cfg, opt.img_size)
 
     # Load weights
-    if opt.weights_path.endswith('.weights'):  # darknet format
-        load_weights(model, opt.weights_path)
-    elif opt.weights_path.endswith('.pt'):  # pytorch format
+    if opt.weights_path.endswith('.pt'):  # pytorch format
         checkpoint = torch.load(opt.weights_path, map_location='cpu')
         model.load_state_dict(checkpoint['model'])
         del checkpoint
+    else:  # darknet format
+        load_weights(model, opt.weights_path)
 
     model.to(device).eval()
 
diff --git a/utils/gcp.sh b/utils/gcp.sh
index f5c2c4d8..3e27c169 100644
--- a/utils/gcp.sh
+++ b/utils/gcp.sh
@@ -13,6 +13,9 @@ python3 detect.py
 # Test
 python3 test.py -img_size 416 -weights_path weights/latest.pt
 
+# Test Darknet
+python3 test.py -img_size 416 -weights_path ../darknet/backup/yolov3.backup
+
 # Download and Test
 sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3
 wget https://pjreddie.com/media/files/yolov3.weights -P weights