update to ultralytics 8.0.35 (#53)

* update model card info

* update to ultralytics 8.0.35

* fix

* Update __init__.py

requirements.txt

@@ -1,5 +1,5 @@
 huggingface_hub>=0.12.0
 fire
-ultralytics>=8.0.25,<8.0.26
+ultralytics>=8.0.35,<8.0.36
 sahi>=0.11.11,<0.12.0
 pandas

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.25"
+__version__ = "0.0.26"

ultralyticsplus/hf_utils.py

@@ -67,15 +67,17 @@ def generate_model_usage_markdown(
         postprocess_str = """print(results[0].probs) # [0.1, 0.2, 0.3, 0.4]
 processed_result = postprocess_classify_output(model, result=results[0])
 print(processed_result) # {"cat": 0.4, "dog": 0.6}"""
-        model_params_str = """model.overrides['conf'] = 0.25  # model confidence threshold"""
+        model_params_str = (
+            """model.overrides['conf'] = 0.25  # model confidence threshold"""
+        )
         metrics_str = f"""      - type: accuracy
         value: {score_top1_acc}  # min: 0.0 - max: 1.0
         name: top1 accuracy
       - type: accuracy
         value: {score_top5_acc}  # min: 0.0 - max: 1.0
         name: top5 accuracy"""
 
-    custom_tags_str = ''
+    custom_tags_str = ""
     if custom_tags:
         if not isinstance(custom_tags, list):
             custom_tags = [custom_tags]
@@ -169,7 +171,7 @@ def generate_thumbnail(
     repo_id=None,
     task="object-detection",
     thumbnail_text=None,
-    export_dir=None
+    export_dir=None,
 ):
     """
     Generate thumbnail for the model card
@@ -248,7 +250,7 @@ def push_model_card_to_hfhub(
     if task in ["object-detection", "instance-segmentation"]:
         sample_visual_path = str(Path(exp_folder) / "val_batch0_labels.jpg")
     elif task == "image-classification":
-        sample_visual_path = 'https://user-images.githubusercontent.com/34196005/212529509-3723ef83-e184-4e57-af37-ed7cfe0faf11.jpg'
+        sample_visual_path = "https://user-images.githubusercontent.com/34196005/212529509-3723ef83-e184-4e57-af37-ed7cfe0faf11.jpg"
     else:
         raise ValueError(f"Task {task} is not supported.")
 
@@ -257,7 +259,7 @@ def push_model_card_to_hfhub(
         repo_id=repo_id,
         task=task,
         thumbnail_text=thumbnail_text,
-        export_dir=exp_folder
+        export_dir=exp_folder,
     )
     upload_file(
         repo_id=repo_id,
@@ -282,7 +284,7 @@ def push_model_card_to_hfhub(
         score_top5_acc=score_top5_acc,
         model_type=model_type,
         labels=labels,
-        custom_tags=custom_tags
+        custom_tags=custom_tags,
     )
     modelcard_path = Path(exp_folder) / "README.md"
     with open(modelcard_path, "w") as file_object:
@@ -308,7 +310,7 @@ def push_config_to_hfhub(
     task="object-detection",
     hf_token=None,
     private=False,
-    model_type="v8"
+    model_type="v8",
 ):
     """
     Pushes a yolov5 config to huggingface hub
@@ -379,8 +381,7 @@ def push_model_to_hfhub(repo_id, exp_folder, hf_token=None, private=False):
     """
     from glob import glob
 
-    from huggingface_hub import (create_repo, delete_file, list_repo_files,
-                                 upload_file)
+    from huggingface_hub import create_repo, delete_file, list_repo_files, upload_file
 
     best_model_path = Path(exp_folder) / "weights/best.pt"
 
@@ -444,7 +445,7 @@ def _push_to_hfhub(
         task=task,
         hf_token=hf_token,
         private=hf_private,
-        model_type=model_type
+        model_type=model_type,
     )
     push_model_card_to_hfhub(
         repo_id=hf_model_id,
@@ -535,18 +536,18 @@ def push_to_hfhub(
         task=task,
         model_type=model.type,
         thumbnail_text=thumbnail_text,
-        custom_tags=custom_tags
+        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"
+            "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",
         }
 
 

ultralyticsplus/ultralytics_utils.py

@@ -5,11 +5,14 @@
 import numpy as np
 from PIL import Image
 from sahi.prediction import ObjectPrediction, PredictionScore
-from sahi.utils.cv import (get_bool_mask_from_coco_segmentation,
-                           read_image_as_pil, visualize_object_predictions)
+from sahi.utils.cv import (
+    get_bool_mask_from_coco_segmentation,
+    read_image_as_pil,
+    visualize_object_predictions,
+)
 from ultralytics import YOLO as YOLOBase
-from ultralytics.nn.tasks import attempt_load_one_weight
-
+from ultralytics.nn.tasks import attempt_load_one_weight, guess_model_task
+from ultralytics.yolo.utils.downloads import GITHUB_ASSET_STEMS
 from ultralyticsplus.hf_utils import download_from_hub
 
 LOGLEVEL = os.environ.get("LOGLEVEL", "INFO").upper()
@@ -42,10 +45,24 @@ def __init__(self, model="yolov8n.yaml", type="v8", hf_token=None) -> None:
         self.overrides = {}  # overrides for trainer object
 
         # Load or create new YOLO model
-        if Path(model).suffix not in (".pt", ".yaml"):
-            self._load_from_hf_hub(model, hf_token=hf_token)
-        else:
-            {".pt": self._load, ".yaml": self._new}[Path(model).suffix](model)
+        suffix = Path(model).suffix
+        if not suffix and Path(model).stem in GITHUB_ASSET_STEMS:
+            model, suffix = (
+                Path(model).with_suffix(".pt"),
+                ".pt",
+            )  # add suffix, i.e. yolov8n -> yolov8n.pt
+        try:
+            if Path(model).suffix not in (".pt", ".yaml"):
+                self._load_from_hf_hub(model, hf_token=hf_token)
+            elif suffix == ".yaml":
+                self._new(model)
+            else:
+                self._load(model)
+        except Exception as e:
+            raise NotImplementedError(
+                f"Unable to load model='{model}'. "
+                f"As an example try model='yolov8n.pt' or model='yolov8n.yaml'"
+            ) from e
 
     def _load_from_hf_hub(self, weights: str, hf_token=None):
         """
@@ -63,8 +80,12 @@ def _load_from_hf_hub(self, weights: str, hf_token=None):
         self.task = self.model.args["task"]
         self.overrides = self.model.args
         self._reset_ckpt_args(self.overrides)
-        self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \
-            self._assign_ops_from_task(self.task)
+        (
+            self.ModelClass,
+            self.TrainerClass,
+            self.ValidatorClass,
+            self.PredictorClass,
+        ) = self._assign_ops_from_task()
 
 
 def render_result(