From 1f5ddf020f6610a0f0525c46bcda24cbdf348d98 Mon Sep 17 00:00:00 2001
From: fatih <34196005+fcakyon@users.noreply.github.com>
Date: Fri, 27 Jan 2023 09:07:15 +0300
Subject: [PATCH] update ultralytics (#46)

* update ultralytics

* add return_dict to push_to_hfhub

* Update __init__.py
---
 requirements.txt            |  2 +-
 ultralyticsplus/__init__.py |  2 +-
 ultralyticsplus/hf_utils.py | 15 +++++++++++++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 7c884a4..59cb78b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 huggingface_hub>=0.12.0
 fire
-ultralytics>=8.0.20,<8.0.21
+ultralytics>=8.0.21,<8.0.22
 sahi>=0.11.11,<0.12.0
 pandas
diff --git a/ultralyticsplus/__init__.py b/ultralyticsplus/__init__.py
index 28f81f0..2875697 100644
--- a/ultralyticsplus/__init__.py
+++ b/ultralyticsplus/__init__.py
@@ -1,4 +1,4 @@
 from .hf_utils import download_from_hub, push_to_hfhub
 from .ultralytics_utils import YOLO, postprocess_classify_output, render_result
 
-__version__ = "0.0.22"
+__version__ = "0.0.23"
diff --git a/ultralyticsplus/hf_utils.py b/ultralyticsplus/hf_utils.py
index 5236e5b..1bdd757 100644
--- a/ultralyticsplus/hf_utils.py
+++ b/ultralyticsplus/hf_utils.py
@@ -4,6 +4,7 @@
 
 import pandas as pd
 from sahi.utils.cv import read_image_as_pil
+
 from ultralyticsplus.other_utils import add_text_to_image
 
 LOGLEVEL = os.environ.get("LOGLEVEL", "INFO").upper()
@@ -24,6 +25,7 @@ def generate_model_usage_markdown(
     custom_tags=None,
 ):
     from ultralytics import __version__ as ultralytics_version
+
     from ultralyticsplus import __version__ as ultralyticsplus_version
 
     hf_task = "image-segmentation" if task == "instance-segmentation" else task
@@ -472,6 +474,7 @@ def push_to_hfhub(
     hf_dataset_id=None,
     thumbnail_text=None,
     custom_tags=None,
+    return_dict=False,
 ):
     """
     Pushes a ultralytics model to huggingface hub
@@ -484,6 +487,7 @@ def push_to_hfhub(
         hf_dataset_id (str): huggingface dataset id to be used for model card
         thumbnail_text (str): text to be used in thumbnail
         custom_tags (list): list of custom tags to be used for model card
+        return_dict (bool): whether to return a dictionary of results or not
     """
     from ultralytics import YOLO
 
@@ -534,6 +538,17 @@ def push_to_hfhub(
         custom_tags=custom_tags
     )
 
+    if return_dict:
+        return {
+            'score_map50': score_map50,
+            'score_map50_mask': score_map50_mask,
+            'score_top1_acc': score_top1_acc,
+            'score_top5_acc': score_top5_acc,
+            'task': task,
+            'model_type': model.type,
+            'thumbnail_url': f"https://huggingface.co/{hf_model_id}/resolve/main/thumbnail.jpg"
+        }
+
 
 def download_from_hub(hf_model_id, hf_token=None):
     """