添加yolo库

yhw-miracle · Jul 2, 2024 · 442fdc2 · 442fdc2
1 parent face0e6
commit 442fdc2
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/.github/workflows/docker-builds.yml b/.github/workflows/docker-builds.yml
@@ -49,5 +49,5 @@ jobs:
           # publish master
           push: ${{ github.ref == 'refs/heads/master' }}
           #tags: acai66/docker_python-opencv-cuda-ffmpeg-gstreamer:${{ matrix.device }}-py${{ matrix.python_version }}-cv${{ matrix.opencv_version }}
-          tags: yhwmiracle/infer:v2.1
-        timeout-minutes: 240
+          tags: yhwmiracle/infer:v3.0
+        timeout-minutes: 240
diff --git a/gpu/Dockerfile b/gpu/Dockerfile
@@ -273,4 +273,8 @@ RUN \
     apt-get update && \
     apt-get install -y gdal-bin python3-gdal
 
+RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+RUN pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
+RUN pip install ultralytics==8.2.1
+
 ENV GST_PLUGIN_FEATURE_RANK="nvcodec:256,nvh264dec:256,nvh265dec:256,nvvp9dec:256,nvvp8dec:256,nvh264enc:256,nvh265enc:256,avdec_h265:128,avdec_h264:128,avdec_vp9:128,avdec_vp8:128,avenc_h264_omx:128"