File size: 1,228 Bytes
3d5c557 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | import json
import os
from typing import List
from numpy import ndarray
from ultralytics.models import YOLO
def import_describe_model() -> YOLO:
current_folder = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
neural_network_file = os.path.join(current_folder, "./yolov8n-fashionpedia-1.torchscript")
return YOLO(neural_network_file, task='detect')
def describe_clothes_batch_opencv_bgr(yolo_model_loaded: YOLO,
pictures_rgb: List[ndarray],
threshold_to_save: float):
results = yolo_model_loaded(pictures_rgb, verbose=False, conf=threshold_to_save)
formatted_results = []
for a_result in results:
formatted_results.append(json.loads(a_result.tojson()))
return formatted_results
def describe_single_clothes_opencv_rgb(yolo_model_loaded: YOLO,
one_clothes_picture_rgb: ndarray,
threshold_to_save: float):
return describe_clothes_batch_opencv_bgr(yolo_model_loaded,
[one_clothes_picture_rgb],
threshold_to_save)[0]
|