Browse Source

Download onnx files from GitHub Release

Kentaro Wada 2 years ago
parent
commit
688ff43300
1 changed files with 13 additions and 2 deletions
  1. 13 2
      labelme/ai/models/segment_anything.py

+ 13 - 2
labelme/ai/models/segment_anything.py

@@ -1,5 +1,6 @@
 import threading
 
+import gdown
 import imgviz
 import numpy as np
 import onnxruntime
@@ -13,10 +14,20 @@ class SegmentAnythingModel:
 
         # encoder_path = "../segment-anything/models/sam_vit_h_4b8939.quantized.encoder.onnx"  # NOQA
         # decoder_path = "../segment-anything/models/sam_vit_h_4b8939.quantized.decoder.onnx"  # NOQA
-        encoder_path = "../segment-anything/models/sam_vit_l_0b3195.quantized.encoder.onnx"  # NOQA
-        decoder_path = "../segment-anything/models/sam_vit_l_0b3195.quantized.decoder.onnx"  # NOQA
+        #
+        # encoder_path = "../segment-anything/models/sam_vit_l_0b3195.quantized.encoder.onnx"  # NOQA
+        # decoder_path = "../segment-anything/models/sam_vit_l_0b3195.quantized.decoder.onnx"  # NOQA
+        #
         # encoder_path = "../segment-anything/models/sam_vit_b_01ec64.quantized.encoder.onnx"  # NOQA
         # decoder_path = "../segment-anything/models/sam_vit_b_01ec64.quantized.decoder.onnx"  # NOQA
+        encoder_path = gdown.cached_download(
+            url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.encoder.onnx",  # NOQA
+            md5="080004dc9992724d360a49399d1ee24b",
+        )
+        decoder_path = gdown.cached_download(
+            url="https://github.com/wkentaro/labelme/releases/download/sam-20230416/sam_vit_l_0b3195.quantized.decoder.onnx",  # NOQA
+            md5="851b7faac91e8e23940ee1294231d5c7",
+        )
 
         self._encoder_session = onnxruntime.InferenceSession(encoder_path)
         self._decoder_session = onnxruntime.InferenceSession(decoder_path)