RB_IBDM_ModelZoo / yolo_utils.py
kiiwee's picture
Upload 5 files
79183dd verified
raw
history blame
513 Bytes
from ultralytics import YOLO
import torch
import numpy as np
import cv2
from huggingface_hub import hf_hub_download
REPO_ID = "kiiwee/Yolov8_InsectDetect"
FILENAME = "insectYolo.pt"
# Ensure you have the model file
model = YOLO(hf_hub_download(repo_id=REPO_ID, filename=FILENAME))
def yolo_processimage(image):
results = model(source=image, show=True,save=True,
conf=0.2, device='mps',save_crop=True)
rgb_image = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
return rgb_image