support task map attribute for yolo >= 8.0.44

tests/test_hf.py

@@ -11,10 +11,16 @@
 hub_id = "ultralyticsplus/yolov8s"
 
 
+# for ultralytics < 8.0.44
 def test_load_from_hub():
     path = download_from_hub(hub_id)
 
 
+# for ultralytics >= 8.0.44
+def test_load_from_hub_yolo_8_0_44():
+    model = YOLO("keremberke/yolov8n-table-extraction")
+
+
 def test_yolo_from_hub():
     model = YOLO(hub_id)
 

ultralyticsplus/ultralytics_utils.py

@@ -83,12 +83,28 @@ 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()
+
+        # for loading model with ultralytics <8.0.44
+        if hasattr(self, "_assign_ops_from_task"):
+            (
+                self.ModelClass,
+                self.TrainerClass,
+                self.ValidatorClass,
+                self.PredictorClass,
+            ) = self.task_map[self.task]
+
+        # for loading model with ultralytics >=8.0.44
+        else:
+            if self.task not in self.task_map:
+                raise ValueError(
+                    f"Task '{self.task}' not supported. Supported tasks: {list(self.task_map.keys())}"
+                )
+            (
+                self.ModelClass,
+                self.TrainerClass,
+                self.ValidatorClass,
+                self.PredictorClass,
+            ) = self.task_map[self.task]
 
 
 def render_result(