Coverage for collectra/tasks/machine_learning/yolo.py: 87.88%

1"""Shared model lifecycle for YOLO-backed tasks.""" 

2 

3from __future__ import annotations 

4 

5from contextlib import contextmanager 

6import importlib 

7from pathlib import Path 

8from typing import TYPE_CHECKING 

9 

10from ...logger import get_logger 

11from ..base import Task 

12 

13if TYPE_CHECKING: 

14 from ultralytics.models import YOLO 

15 

16__all__ = ["YOLOTask"] 

17 

18logger = get_logger(__name__) 

19 

20 

21def get_yolo_class(): 

22 """Import the optional backend when a YOLO task executes.""" 

23 try: 

24 from ultralytics.models import YOLO 

25 except ModuleNotFoundError as error: 

26 if error.name != "ultralytics": 

27 raise 

28 raise ImportError( 

29 'YOLO tasks require optional dependencies. Install with: ' 

30 'pip install "collectra[yolo]"' 

31 ) from error 

32 return YOLO 

33 

34 

35class YOLOTask(Task): 

36 """Base class for tasks backed by an Ultralytics YOLO model.""" 

37 

38 model: str | Path | YOLO 

39 original_model_path: str | Path = "" 

40 

41 def __init__(self, name: str, model: str | Path | YOLO = "", **kwargs) -> None: 

42 super().__init__(name, model=model, **kwargs) 

43 

44 @contextmanager 

45 def _wandb_logging(self, enabled: bool = False): 

46 """Select Ultralytics' native W&B integration for this training run.""" 

47 from ultralytics import settings 

48 

49 previous = settings["wandb"] 

50 module_name = "ultralytics.utils.callbacks.wb" 

51 callbacks = getattr(self.model, "callbacks", {}) 

52 

53 def remove_wandb_callbacks(): 

54 # Reused models can retain callbacks from an earlier enabled run. 

55 for handlers in callbacks.values(): 

56 handlers[:] = [ 

57 handler 

58 for handler in handlers 

59 if getattr(handler, "__module__", None) != module_name 

60 ] 

61 

62 integration = None 

63 try: 

64 if previous != enabled: 

65 settings.update({"wandb": enabled}) 

66 # The native integration checks settings at import time. Reload it 

67 # so enabling/disabling logging also works on subsequent runs. 

68 integration = importlib.import_module(module_name) 

69 importlib.reload(integration) 

70 remove_wandb_callbacks() 

71 yield 

72 finally: 

73 remove_wandb_callbacks() 

74 if settings["wandb"] != previous: 

75 settings.update({"wandb": previous}) 

76 if integration is not None: 

77 importlib.reload(integration) 

78 

79 def _init_model(self) -> None: 

80 """Load and validate the configured YOLO model.""" 

81 YOLO = get_yolo_class() 

82 

83 self._load() 

84 if not isinstance(self.model, YOLO): 

85 raise ValueError("Model must be a YOLO instance") 

86 

87 def _reload(self) -> None: 

88 """Reload the model from its original path.""" 

89 if not self.original_model_path: 

90 logger.warning("Original model path is not set. Skipping...") 

91 return 

92 self.model = self.original_model_path 

93 self._load() 

94 

95 def _load(self) -> None: 

96 """Load a YOLO model only when execution or training requires it.""" 

97 YOLO = get_yolo_class() 

98 

99 if self.model and isinstance(self.model, (str, Path)): 

100 self.original_model_path = self.model 

101 self.model = YOLO(Path(self.model)) 

102 config = getattr(getattr(self.model, "model", None), "yaml", None) 

103 if isinstance(config, dict) and "feature_fusion" in config: 

104 from ..object_detection.yolo_feature_fusion import enable_feature_fusion 

105 

106 enable_feature_fusion(self.model) 

107 if not isinstance(self.model, YOLO): 

108 logger.warning("Invalid YOLO model value: %r", self.model)