diff --git a/Dockerfile b/Dockerfile index b771e9ed..89bc72ab 100644 --- a/Dockerfile +++ b/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 + diff --git a/models.py b/models.py index abca7e61..e3177dad 100755 --- a/models.py +++ b/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 diff --git a/utils/google_utils.py b/utils/google_utils.py index 211d5485..303265b4 100644 --- a/utils/google_utils.py +++ b/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'):