Coverage for collectra/tasks/machine_learning/training.py: 82.96%

1"""Opt-in utilities for training task implementations.""" 

2 

3from __future__ import annotations 

4 

5import shutil 

6from datetime import datetime 

7from pathlib import Path 

8from typing import Callable 

9 

10import yaml 

11from rich.console import Console 

12from rich.table import Table 

13 

14from collectra.types.base import ArtefactNode 

15from collectra.types.images import Image, ImageCrop 

16from collectra.utils import change_dir 

17 

18from ...logger import get_logger 

19 

20logger = get_logger(__name__) 

21 

22 

23def print_distribution_table( 

24 title: str, 

25 classes: list[str], 

26 train_counts: dict[str, int], 

27 validation_counts: dict[str, int], 

28) -> None: 

29 """Print class counts and each class's held-out validation percentage.""" 

30 table = Table(title=title, show_lines=True) 

31 table.add_column( 

32 "Class Name", justify="left", style="green", header_style="bold green" 

33 ) 

34 table.add_column( 

35 "Train Count", justify="right", style="red", header_style="bold red" 

36 ) 

37 table.add_column( 

38 "Validation Count", justify="right", style="blue", header_style="bold blue" 

39 ) 

40 table.add_column("Validation %", justify="right", header_style="bold") 

41 

42 total_train = 0 

43 total_validation = 0 

44 for class_name in classes: 

45 train_count = train_counts[class_name] 

46 validation_count = validation_counts[class_name] 

47 total = train_count + validation_count 

48 percentage = 100 * validation_count / total if total else 0.0 

49 table.add_row( 

50 class_name, 

51 str(train_count), 

52 str(validation_count), 

53 f"{percentage:.1f}%", 

54 ) 

55 total_train += train_count 

56 total_validation += validation_count 

57 

58 grand_total = total_train + total_validation 

59 total_percentage = 100 * total_validation / grand_total if grand_total else 0.0 

60 table.add_row( 

61 "Total", 

62 str(total_train), 

63 str(total_validation), 

64 f"{total_percentage:.1f}%", 

65 style="bold", 

66 ) 

67 Console().print(table) 

68 

69 

70def prepare_object_detection_inputs( 

71 processed_inputs: list, 

72 processed_parents: list, 

73 input_maps: dict[str, list], 

74 parent_input_maps: dict[str, list], 

75) -> list: 

76 """Attach source images to detection crops and return the task inputs.""" 

77 for item in processed_inputs: 

78 if type(item) is Image: 

79 continue 

80 if type(item) is not ImageCrop: 

81 raise TypeError(f"Input {item.id} is not an ImageCrop") 

82 if isinstance(item.parents, str): 

83 parent_id = item.parents 

84 elif isinstance(item.parents, list) and len(item.parents) == 1: 

85 parent_id = item.parents[0] 

86 else: 

87 continue 

88 parent = next( 

89 ( 

90 candidate 

91 for candidate in processed_parents 

92 if candidate.id == parent_id and item.data == candidate.data 

93 ), 

94 None, 

95 ) 

96 if not isinstance(parent, Image): 

97 raise ValueError( 

98 f"Image parent {parent_id!r} for input {item.id!r} was not found" 

99 ) 

100 item.add_source_parent(parent) 

101 return processed_inputs 

102 

103 

104def _training_files(task, inputs: list[str]) -> list[Path]: 

105 pipeline = _pipeline_for(task) 

106 files = [] 

107 for value in inputs: 

108 path = Path(value) 

109 if path.is_dir(): 

110 files.extend(path.glob(f"*.{pipeline.ext}")) 

111 elif path.suffix == f".{pipeline.ext}": 

112 files.append(path) 

113 return files 

114 

115 

116def _file_partition(item_file: Path) -> str: 

117 result_file = item_file / "results.yaml" 

118 if not result_file.exists(): 

119 return "" 

120 with open(result_file, "r") as f: 

121 file_data = yaml.safe_load(f) or {} 

122 metadata = file_data.get("collectra_results_metadata") or {} 

123 return metadata.get("partition") or "" 

124 

125 

126def _limit_training_files( 

127 files: list[Path], max_items: int, validation: str = "", exclude: str = "" 

128) -> list[Path]: 

129 """Keep the first ``max_items`` training and ``max_items`` validation files. 

130 

131 Only the partition metadata is read, and reading stops once both 

132 quotas are filled, so large datasets are not loaded for quick tests. 

133 """ 

134 selected = [] 

135 counts = {"train": 0, "validation": 0} 

136 for item_file in sorted(files): 

137 partition = _file_partition(item_file) 

138 if exclude and partition == exclude: 

139 continue 

140 split = "validation" if validation and partition == validation else "train" 

141 if counts[split] < max_items: 

142 selected.append(item_file) 

143 counts[split] += 1 

144 if counts["train"] >= max_items and ( 

145 not validation or counts["validation"] >= max_items 

146 ): 

147 break 

148 logger.info( 

149 "max_items=%d: using %d training and %d validation files", 

150 max_items, 

151 counts["train"], 

152 counts["validation"], 

153 ) 

154 return selected 

155 

156 

157def _prepare_data( 

158 task, 

159 files: list[Path], 

160 nodes: list[ArtefactNode], 

161 *, 

162 include_unlabelled: bool = False, 

163): 

164 processed = [] 

165 input_maps = {} 

166 for item_file in files: 

167 input_maps[item_file.name] = ArtefactNode.batch_process( 

168 item_file, 

169 nodes, 

170 include_unlabelled=include_unlabelled, 

171 ) 

172 processed.extend(input_maps[item_file.name]) 

173 return processed, input_maps 

174 

175 

176def _pipeline_for(task): 

177 if task.pipeline is None: 

178 raise RuntimeError(f"Task {task.name!r} is not attached to a pipeline") 

179 return task.pipeline 

180 

181 

182def save_training_result(task, results, *, log: str, **kwargs) -> None: 

183 """Copy the best model and optional class metadata into the pipeline.""" 

184 pipeline = _pipeline_for(task) 

185 best_model_path = results.save_dir / "weights" / "best.pt" 

186 if not best_model_path.exists(): 

187 raise FileNotFoundError(f"Best model file not found at {best_model_path}") 

188 logger.info("Training metrics: %s", results.results_dict) 

189 current_model = pipeline.data.get(task.name, {}).get("model", "") 

190 if best_model_path.name == current_model: 

191 return 

192 new_model_path = f"{task.name}-{Path(log).name}.pt" 

193 shutil.copy(best_model_path, new_model_path) 

194 pipeline.data[task.name]["model"] = new_model_path 

195 classes_src = best_model_path.parent / "classes.json" 

196 if classes_src.exists(): 

197 shutil.copy(classes_src, f"{Path(new_model_path).stem}.classes.json") 

198 

199 

200def train_from_files( 

201 task, 

202 inputs: list[str], 

203 trainer: Callable, 

204 *, 

205 prepare_inputs: Callable | None = None, 

206 include_unlabelled: bool = False, 

207 **kwargs, 

208): 

209 """Resolve pipeline artefacts and invoke a concrete backend trainer.""" 

210 pipeline = _pipeline_for(task) 

211 task_node = pipeline.node_manager.resolve_node(task.name) 

212 children = pipeline.node_manager.get_children_artefact(task_node) 

213 parents = pipeline.node_manager.get_ancestor_artefacts(task_node) 

214 kwargs["classes"] = kwargs.get("classes", [child.name for child in children]) 

215 kwargs = pipeline._merge_task_params(task.name, kwargs) 

216 files = _training_files(task, inputs) 

217 max_items = int(kwargs.get("max_items", 0) or 0) 

218 if max_items > 0: 

219 files = _limit_training_files( 

220 files, 

221 max_items, 

222 validation=kwargs.get("validation", ""), 

223 exclude=kwargs.get("exclude", ""), 

224 ) 

225 processed_inputs, input_maps = _prepare_data( 

226 task, 

227 files, 

228 children, 

229 include_unlabelled=include_unlabelled, 

230 ) 

231 processed_parents, parent_input_maps = _prepare_data(task, files, parents) 

232 if prepare_inputs is not None: 

233 processed_inputs = prepare_inputs( 

234 processed_inputs, 

235 processed_parents, 

236 input_maps, 

237 parent_input_maps, 

238 ) 

239 with change_dir(pipeline.path): 

240 results = trainer(*processed_inputs, **kwargs) 

241 save_training_result(task, results, **kwargs) 

242 return results 

243 

244 

245def run_training_command( 

246 task, 

247 inputs: list[str], 

248 trainer: Callable, 

249 *, 

250 keep_log: bool, 

251 output: Path | None, 

252 validation: str, 

253 exclude: str, 

254 prepare_inputs: Callable | None = None, 

255 include_unlabelled: bool = False, 

256 **kwargs, 

257): 

258 """Run the common filesystem lifecycle requested by a concrete command.""" 

259 pipeline = _pipeline_for(task) 

260 log = ( 

261 Path(output).expanduser().resolve() 

262 if output is not None 

263 else Path.cwd() / datetime.now().strftime("%Y%m%d_%H%M%S") 

264 ) 

265 results = train_from_files( 

266 task, 

267 inputs, 

268 trainer, 

269 prepare_inputs=prepare_inputs, 

270 include_unlabelled=include_unlabelled, 

271 project=f"{pipeline.path.name}-{task.name}", 

272 log=log, 

273 base_folder=Path.cwd(), 

274 validation=validation, 

275 exclude=exclude, 

276 **kwargs, 

277 ) 

278 pipeline.save(task.name) 

279 if not keep_log: 

280 shutil.rmtree(log, ignore_errors=True) 

281 return results