updates

Dockerfile

@@ -12,11 +12,7 @@ FROM nvcr.io/nvidia/pytorch:19.08-py3
 #WORKDIR /usr/src
 #RUN pip uninstall -y opencv-python
 #RUN apt-get update
-#RUN apt-get install -y gstreamer1.0-python3-dbg-plugin-loader
-## RUN apt-get install gstreamer1.0
-## RUN apt install -y ubuntu-restricted-extras
-#RUN apt install -y libgstreamer1.0-dev
-#RUN apt install -y libgstreamer-plugins-base1.0-dev
+#RUN apt-get install -y gstreamer1.0-python3-dbg-plugin-loader libgstreamer1.0-dev libgstreamer-plugins-base1.0-dev
 #RUN git clone https://github.com/opencv/opencv.git && cd opencv && git checkout 4.1.1 && mkdir build
 #RUN git clone https://github.com/opencv/opencv_contrib.git && cd opencv_contrib && git checkout 4.1.1
 #RUN cd opencv/build && cmake ../ \
@@ -26,7 +22,7 @@ FROM nvcr.io/nvidia/pytorch:19.08-py3
 #    -D PYTHON3_INCLUDE_PATH=/opt/conda/include/python3.6m \
 #    -D PYTHON3_LIBRARIES=/opt/conda/lib/python3.6/site-packages \
 #    -D WITH_GSTREAMER=ON \
-#    -D WITH_FFMPEG=OFF \
+#    -D WITH_FFMPEG=ON \
 #    && make && make install && ldconfig
 #RUN cd /usr/local/lib/python3.6/site-packages/cv2/python-3.6/ && mv cv2.cpython-36m-x86_64-linux-gnu.so cv2.so
 #RUN cd /opt/conda/lib/python3.6/site-packages/ && ln -s /usr/local/lib/python3.6/site-packages/cv2/python-3.6/cv2.so cv2.so
@@ -58,3 +54,4 @@ COPY . /usr/src/app
 
 # Build and Push
 # export tag=ultralytics/yolov3:v0 && sudo docker build -t $tag . && docker push $tag
+

models.py

@@ -2,6 +2,7 @@ import torch.nn.functional as F
 
 from utils.parse_config import *
 from utils.utils import *
+from utils.google_utils import *
 
 ONNX_EXPORT = False
 
@@ -296,13 +297,18 @@ def load_darknet_weights(self, weights, cutoff=-1):
     # Try to download weights if not available locally
     msg = weights + ' missing, download from https://drive.google.com/drive/folders/1uxgUBemJVw9wZsdpboYbzUN4bcRhsuAI'
     if not os.path.isfile(weights):
-        try:
-            url = 'https://pjreddie.com/media/files/' + file
-            print('Downloading ' + url)
-            os.system('curl -f ' + url + ' -o ' + weights)
-        except IOError:
-            print(msg)
-            os.system('rm ' + weights)  # remove partial downloads
+        if file == 'yolov3-spp.weights':
+            gdrive_download(id='1oPCHKsM2JpM-zgyepQciGli9X0MTsJCO', name=weights)
+        elif file == 'darknet.53.conv.74':
+            gdrive_download(id='18xqvs_uwAqfTXp-LJCYLYNHBOcrwbrp0', name=weights)
+        else:
+            try:  # download from pjreddie.com
+                url = 'https://pjreddie.com/media/files/' + file
+                print('Downloading ' + url)
+                os.system('curl -f ' + url + ' -o ' + weights)
+            except IOError:
+                print(msg)
+                os.system('rm ' + weights)  # remove partial downloads
     assert os.path.exists(weights), msg  # download missing weights from Google Drive
 
     # Establish cutoffs

utils/google_utils.py

@@ -18,17 +18,17 @@ def gdrive_download(id='1HaXkef9z6y5l4vUnCYgdmEAj61c6bfWO', name='coco.zip'):
     if os.path.exists(name):  # remove existing
         os.remove(name)
 
-    # Attempt small file download
-    s = 'curl -f -L -o %s https://drive.google.com/uc?export=download&id=%s' % (name, id)
-    os.system(s)
-
     # Attempt large file download
-    if not os.path.exists(name):  # file size > 40MB
-        s = ["curl -c ./cookie -s -L \"https://drive.google.com/uc?export=download&id=%s\" > /dev/null" % id,
-             "curl -Lb ./cookie \"https://drive.google.com/uc?export=download&confirm=`awk '/download/ {print $NF}' ./cookie`&id=%s\" -o %s" % (
-                 id, name),
-             'rm ./cookie']
-        [os.system(x) for x in s]  # run commands
+    s = ["curl -c ./cookie -s -L \"https://drive.google.com/uc?export=download&id=%s\" > /dev/null" % id,
+         "curl -Lb ./cookie -s \"https://drive.google.com/uc?export=download&confirm=`awk '/download/ {print $NF}' ./cookie`&id=%s\" -o %s" % (
+             id, name),
+         'rm ./cookie']
+    [os.system(x) for x in s]  # run commands
+
+    # Attempt small file download
+    if not os.path.exists(name):  # file size < 40MB
+        s = 'curl -f -L -o %s https://drive.google.com/uc?export=download&id=%s' % (name, id)
+        os.system(s)
 
     # Unzip if archive
     if name.endswith('.zip'):