Update Lib/BirdNestDetection.py

Lib/BirdNestDetection.py

@@ -9,6 +9,8 @@
 """
 import os
 import sys
+
+
 sys.path.append(os.getcwd())
 
 import random
@@ -18,9 +20,17 @@ import torchvision
 import onnxruntime as ort
 import cv2
 import numpy as np
+
+from Tool.Core import DownloadHFModel
 from Lib.Const import LABELS, COLOR_MAP, COLOR_MAP_RGB
 
 
+# Download Model
+REPO_ID = "blitzkrieg0000/yolov7_bird-nest-detection"
+MODEL_FILE = "yolov7_kus.onnx"
+DownloadHFModel(REPO_ID, MODEL_FILE)
+
+
 def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
     # Rescale coords (xyxy) from img1_shape to img0_shape
     if ratio_pad is None:  # calculate from img0_shape
@@ -223,7 +233,7 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=3):
 
 
 print(ort.get_available_providers())
-session = ort.InferenceSession("Weight/yolov7_kus.onnx", providers=ort.get_available_providers())
+session = ort.InferenceSession(os.path.join("./Weight", MODEL_FILE), providers=ort.get_available_providers())
 
 input_name = session.get_inputs()[0].name
 print("input name", input_name)