Coverage for collectra/tasks/object_detection/yolo.py: 76.08%

1from __future__ import annotations 

2 

3import shutil 

4import tempfile 

5from collections import Counter 

6from pathlib import Path 

7from typing import Annotated, TYPE_CHECKING 

8 

9from cappa import Arg 

10 

11from collectra.cli import command 

12from .adaptation_commands import YOLOAdaptationCommands 

13from collectra.types.images import ( 

14 Image, 

15 ImageCrop, 

16 image_channel_count, 

17 read_tiff_channels, 

18) 

19from collectra.types.texts import Text 

20from collectra.utils import change_dir 

21 

22from collectra.utils import threading_locked 

23from ..machine_learning.training import ( 

24 prepare_object_detection_inputs, 

25 print_distribution_table, 

26 run_training_command, 

27) 

28from ..machine_learning.yolo import YOLOTask, get_yolo_class 

29 

30if TYPE_CHECKING: 

31 from ultralytics.engine.results import Results 

32 from ultralytics.models import YOLO 

33 from ultralytics.utils.metrics import ClassifyMetrics, DetMetrics 

34 

35__all__ = ["ObjectDetectionYOLO"] 

36 

37 

38class DetectionResult: 

39 metrics: dict 

40 best_result: DetMetrics 

41 

42 def __init__(self, result: DetMetrics) -> None: 

43 self.best_result = result 

44 self.metrics = result.results_dict if result else {} 

45 

46 @staticmethod 

47 def get_best_result( 

48 results: list["DetectionResult"], metric: str = "metrics/mAP50-95(B)" 

49 ) -> DetMetrics: 

50 best_result: DetMetrics | None = None 

51 best_metric_result = 0.0 

52 for result in results: 

53 if result is None: 

54 continue 

55 metric_result = result.metrics.get(metric, 0.0) 

56 if best_result is None: 

57 best_metric_result = metric_result 

58 best_result = result.best_result 

59 continue 

60 if metric_result > best_metric_result: 

61 best_metric_result = metric_result 

62 best_result = result.best_result 

63 if best_result is None: 

64 raise ValueError("No valid results found to determine the best result.") 

65 return best_result 

66 

67 

68class ObjectDetectionYOLO(YOLOAdaptationCommands, YOLOTask): 

69 singletons: bool = False 

70 imgsz: int | None = None 

71 

72 @command(name="run") 

73 def cli_run( 

74 self, 

75 inputs: list[str], 

76 output: Path | None = None, 

77 verbose: bool = False, 

78 usage: bool = False, 

79 render: bool = False, 

80 model: Annotated[ 

81 Path | None, 

82 Arg(help="Model checkpoint path. Omit to keep the configured model."), 

83 ] = None, 

84 imgsz: Annotated[ 

85 int | None, 

86 Arg( 

87 help="Inference image size in pixels. Omit to keep the model's setting." 

88 ), 

89 ] = None, 

90 ) -> None: 

91 """Run this object detector for the supplied input files.""" 

92 if model is not None: 

93 self.model = model 

94 if imgsz is not None: 

95 self.imgsz = imgsz 

96 super().cli_run( 

97 inputs, output=output, verbose=verbose, usage=usage, render=render 

98 ) 

99 

100 @threading_locked() 

101 def run( 

102 self, 

103 *args: Image, 

104 model: str | Path | None = None, 

105 imgsz: int | None = None, 

106 **kwargs, 

107 ) -> list[ImageCrop]: 

108 """Run object detection inference on the provided Image. 

109 

110 This method performs object detection on the provided input image 

111 using the trained YOLO model. 

112 

113 Results include bounding boxes, confidence scores, and class 

114 predictions for each detected object. 

115 

116 Args: 

117 input (Image): The input image on which to perform object detection. 

118 model: Optional checkpoint path overriding the task's configured model. 

119 imgsz: Optional inference image size in pixels. Defaults to the task's 

120 configured size, or the model's setting when unset. 

121 

122 Returns: 

123 list[ImageCrop]: A list of ImageCrop objects representing the detected 

124 objects in the input image. 

125 """ 

126 if len(args) != 1: 

127 raise ValueError("This task only supports a single Image input.") 

128 if not isinstance(args[0], Image): 

129 raise TypeError("Input must be an instance of Image.") 

130 image: Image = args[0] 

131 if model is not None: 

132 self.model = model 

133 self._init_model() 

134 from torchvision.ops import batched_nms 

135 

136 inference_kwargs = {"iou": 0.7, "conf": 0.25} 

137 imgsz = imgsz if imgsz is not None else self.imgsz 

138 if imgsz is not None: 

139 inference_kwargs["imgsz"] = imgsz 

140 

141 if image.get_path().suffix.lower() in {".tif", ".tiff"}: 

142 import numpy as np 

143 import torch 

144 

145 pixels = read_tiff_channels(image.get_path()) 

146 if isinstance(image, ImageCrop): 

147 left, top, right, bottom = image.coordinates() 

148 if right <= left or bottom <= top: 

149 raise ValueError( 

150 f"ImageCrop {image.id!r} produces an empty pixel crop" 

151 ) 

152 pixels = pixels[top:bottom, left:right, :] 

153 pixels = np.rot90(pixels, k=image.orientation.to_degree() // 90) 

154 first_conv = next( 

155 layer 

156 for layer in self.model.model.modules() 

157 if isinstance(layer, torch.nn.Conv2d) 

158 ) 

159 expected_channels = first_conv.in_channels 

160 model_config = getattr(self.model.model, "yaml", {}) 

161 if "feature_fusion" in model_config: 

162 expected_channels = model_config["channels"] 

163 if pixels.shape[2] != expected_channels: 

164 raise ValueError( 

165 f"{image.get_path()}: {pixels.shape[2]} image channels, " 

166 f"but the detection model expects {expected_channels}" 

167 ) 

168 # Ultralytics expects BGR for three-channel arrays, but preserves 

169 # channel order for multispectral arrays. Keep its usual resizing, 

170 # normalization, and box scaling relative to the original image. 

171 if pixels.shape[2] == 3: 

172 pixels = pixels[..., ::-1] 

173 results: Results = self.model( 

174 np.ascontiguousarray(pixels), **inference_kwargs 

175 )[0] 

176 elif isinstance(image, ImageCrop): 

177 with tempfile.TemporaryDirectory() as temp_dir: 

178 img = image.pil() 

179 img_path = Path(temp_dir) / f"{Path(image.get_path()).stem}.png" 

180 img.save(img_path, format="PNG") 

181 results: Results = (self.model(img_path, **inference_kwargs))[0] 

182 else: 

183 results: Results = (self.model(image.get_path(), **inference_kwargs))[0] 

184 

185 if getattr(self.model.model, "end2end", False) and len(results.boxes): 

186 keep = batched_nms( 

187 results.boxes.xyxy, 

188 results.boxes.conf, 

189 results.boxes.cls, 

190 iou_threshold=0.7, 

191 ) 

192 results = results[keep] 

193 

194 detections: list[ImageCrop] = [] 

195 

196 if results.boxes is None or len(results.boxes) == 0: 

197 return detections 

198 

199 coordinates = results.boxes.xywhn.clone() 

200 names = [ 

201 results.names[class_name.int().item()] for class_name in results.boxes.cls 

202 ] 

203 

204 counts = Counter() 

205 if len(names) == 0: 

206 # No objects detected, return empty list 

207 return detections 

208 for index in range(len(coordinates)): 

209 name = names[index] 

210 

211 if self.singletons and counts[name]: 

212 continue 

213 

214 x, y, w, h = coordinates[index] 

215 confidence = results.boxes.conf[index] 

216 

217 image_crop = image.make_crop( 

218 x_center=float(x), 

219 y_center=float(y), 

220 width_relative=float(w), 

221 height_relative=float(h), 

222 confidence=confidence.item(), 

223 orientation=image.orientation, 

224 name=name, 

225 ) 

226 counts[name] += 1 

227 detections.append(image_crop) 

228 

229 print(f"Found {len(detections)} objects in the image:") 

230 for name, count in counts.items(): 

231 print(f"\t{name}: {count}") 

232 return detections 

233 

234 @command 

235 def train( 

236 self, 

237 inputs: Annotated[ 

238 list[str], 

239 Arg( 

240 help="Pipeline data files or directories containing training annotations." 

241 ), 

242 ], 

243 model: Annotated[ 

244 str, 

245 Arg( 

246 help="YOLO weights name or checkpoint path (for example yolo26n.pt), or a model YAML file. Omit to keep the configured model." 

247 ), 

248 ] = "", 

249 output: Annotated[ 

250 Path | None, 

251 Arg(help="Directory for training logs, dataset files, and checkpoints."), 

252 ] = None, 

253 keep_log: Annotated[ 

254 bool, 

255 Arg( 

256 help="Keep the training directory after saving the model to the pipeline." 

257 ), 

258 ] = True, 

259 validation: Annotated[ 

260 str, 

261 Arg(help="Partition value identifying validation examples."), 

262 ] = "", 

263 exclude: Annotated[ 

264 str, 

265 Arg(help="Partition value identifying examples to exclude from training."), 

266 ] = "", 

267 epochs: Annotated[ 

268 int, 

269 Arg(help="Number of training epochs."), 

270 ] = 100, 

271 batch: Annotated[ 

272 int, 

273 Arg(help="Training batch size."), 

274 ] = 16, 

275 imgsz: Annotated[ 

276 int, 

277 Arg(help="Target training image size in pixels."), 

278 ] = 640, 

279 early_stop: Annotated[ 

280 int, 

281 Arg( 

282 help="Stop after this many epochs without validation improvement; 0 disables early stopping." 

283 ), 

284 ] = 50, 

285 preview: Annotated[ 

286 bool, 

287 Arg( 

288 help="Preview the prepared training images and annotations before training." 

289 ), 

290 ] = False, 

291 workers: Annotated[ 

292 int, 

293 Arg(help="Number of data-loading workers per distributed process."), 

294 ] = 8, 

295 pretrained: Annotated[ 

296 bool, 

297 Arg(help="Use pretrained weights when supported by the selected model."), 

298 ] = True, 

299 optimizer: Annotated[ 

300 str, 

301 Arg( 

302 help="Optimizer: auto, SGD, MuSGD, Adam, Adamax, AdamW, NAdam, RAdam, or RMSProp." 

303 ), 

304 ] = "auto", 

305 seed: Annotated[ 

306 int, 

307 Arg(help="Random seed for training."), 

308 ] = 0, 

309 deterministic: Annotated[ 

310 bool, 

311 Arg( 

312 help="Use deterministic operations for reproducibility; may reduce speed." 

313 ), 

314 ] = True, 

315 single_cls: Annotated[ 

316 bool, 

317 Arg(help="Treat all object classes as one class."), 

318 ] = False, 

319 rect: Annotated[ 

320 bool, 

321 Arg(help="Use rectangular training batches to reduce padding."), 

322 ] = False, 

323 cos_lr: Annotated[ 

324 bool, 

325 Arg(help="Use a cosine learning-rate schedule."), 

326 ] = False, 

327 close_mosaic: Annotated[ 

328 int, 

329 Arg( 

330 help="Disable mosaic augmentation for the final N epochs; 0 keeps it enabled." 

331 ), 

332 ] = 10, 

333 amp: Annotated[ 

334 bool, 

335 Arg(help="Enable automatic mixed precision training."), 

336 ] = True, 

337 fraction: Annotated[ 

338 float, 

339 Arg(help="Fraction of the training dataset to use; 1.0 uses all examples."), 

340 ] = 1.0, 

341 freeze: Annotated[ 

342 int | None, 

343 Arg(help="Freeze the first N model layers during training."), 

344 ] = None, 

345 multi_scale: Annotated[ 

346 float, 

347 Arg( 

348 help="Vary image size by this fraction of imgsz; 0 disables multi-scale training." 

349 ), 

350 ] = 0.0, 

351 lr0: Annotated[ 

352 float, 

353 Arg(help="Initial learning rate."), 

354 ] = 0.01, 

355 lrf: Annotated[ 

356 float, 

357 Arg(help="Final learning-rate multiplier; final rate is lr0 times lrf."), 

358 ] = 0.01, 

359 momentum: Annotated[ 

360 float, 

361 Arg(help="SGD momentum or the first Adam beta coefficient."), 

362 ] = 0.937, 

363 weight_decay: Annotated[ 

364 float, 

365 Arg(help="Weight decay used for regularization."), 

366 ] = 0.0005, 

367 warmup_epochs: Annotated[ 

368 float, 

369 Arg( 

370 help="Number of learning-rate warmup epochs; fractional values are allowed." 

371 ), 

372 ] = 3.0, 

373 warmup_momentum: Annotated[ 

374 float, 

375 Arg(help="Initial momentum during warmup."), 

376 ] = 0.8, 

377 warmup_bias_lr: Annotated[ 

378 float, 

379 Arg(help="Initial learning rate for bias parameters during warmup."), 

380 ] = 0.1, 

381 box: Annotated[ 

382 float, 

383 Arg(help="Weight applied to bounding-box loss."), 

384 ] = 7.5, 

385 cls: Annotated[ 

386 float, 

387 Arg(help="Weight applied to classification loss."), 

388 ] = 0.5, 

389 cls_pw: Annotated[ 

390 float, 

391 Arg( 

392 help="Class-imbalance weighting power; 0 disables weighting, 1 uses inverse class frequency." 

393 ), 

394 ] = 0.0, 

395 dfl: Annotated[ 

396 float, 

397 Arg(help="Weight applied to distribution focal loss."), 

398 ] = 1.5, 

399 nbs: Annotated[ 

400 int, 

401 Arg(help="Nominal batch size used for loss normalization."), 

402 ] = 64, 

403 hsv_h: Annotated[ 

404 float, 

405 Arg(help="Hue augmentation amount as a fraction."), 

406 ] = 0.015, 

407 hsv_s: Annotated[ 

408 float, 

409 Arg(help="Saturation augmentation amount as a fraction."), 

410 ] = 0.7, 

411 hsv_v: Annotated[ 

412 float, 

413 Arg(help="Brightness augmentation amount as a fraction."), 

414 ] = 0.4, 

415 degrees: Annotated[ 

416 float, 

417 Arg(help="Maximum random rotation in either direction, in degrees."), 

418 ] = 0.0, 

419 translate: Annotated[ 

420 float, 

421 Arg(help="Maximum random translation as a fraction of image size."), 

422 ] = 0.1, 

423 scale: Annotated[ 

424 float, 

425 Arg(help="Random scale variation around 1.0."), 

426 ] = 0.5, 

427 shear: Annotated[ 

428 float, 

429 Arg(help="Maximum random shear in either direction, in degrees."), 

430 ] = 0.0, 

431 perspective: Annotated[ 

432 float, 

433 Arg(help="Strength of random perspective augmentation."), 

434 ] = 0.0, 

435 flipud: Annotated[ 

436 float, 

437 Arg(help="Probability of vertically flipping an image."), 

438 ] = 0.0, 

439 fliplr: Annotated[ 

440 float, 

441 Arg(help="Probability of horizontally flipping an image."), 

442 ] = 0.5, 

443 bgr: Annotated[ 

444 float, 

445 Arg(help="Probability of swapping RGB and BGR color channels."), 

446 ] = 0.0, 

447 mosaic: Annotated[ 

448 float, 

449 Arg(help="Probability of mosaic augmentation."), 

450 ] = 1.0, 

451 mixup: Annotated[ 

452 float, 

453 Arg(help="Probability of MixUp augmentation."), 

454 ] = 0.0, 

455 cutmix: Annotated[ 

456 float, 

457 Arg(help="Probability of CutMix augmentation."), 

458 ] = 0.0, 

459 time: Annotated[ 

460 float, 

461 Arg( 

462 help="Training time limit in hours; a positive value overrides epochs, 0 leaves it unset." 

463 ), 

464 ] = 0.0, 

465 save: Annotated[ 

466 bool, 

467 Arg(help="Save training checkpoints."), 

468 ] = True, 

469 save_period: Annotated[ 

470 int, 

471 Arg( 

472 help="Save an additional checkpoint every N epochs; -1 disables periodic saves." 

473 ), 

474 ] = -1, 

475 cache: Annotated[ 

476 str, 

477 Arg(help="Cache images in ram or on disk; omit to leave caching disabled."), 

478 ] = "", 

479 device: Annotated[ 

480 str, 

481 Arg( 

482 help="Training device, such as cpu, mps, 0, or 0,1; omit for automatic selection." 

483 ), 

484 ] = "", 

485 verbose: Annotated[ 

486 bool, 

487 Arg(help="Print detailed training logs."), 

488 ] = True, 

489 resume: Annotated[ 

490 bool, 

491 Arg(help="Resume training from the selected checkpoint."), 

492 ] = False, 

493 compile: Annotated[ 

494 str, 

495 Arg( 

496 help="Torch compilation mode: default, reduce-overhead, or max-autotune-no-cudagraphs; omit to disable." 

497 ), 

498 ] = "", 

499 distill_model: Annotated[ 

500 str, 

501 Arg( 

502 help="Teacher model checkpoint for knowledge distillation; omit to disable." 

503 ), 

504 ] = "", 

505 dis: Annotated[ 

506 float, 

507 Arg(help="Weight applied to knowledge-distillation loss."), 

508 ] = 6.0, 

509 plots: Annotated[ 

510 bool, 

511 Arg(help="Save training and validation plots and images."), 

512 ] = True, 

513 wandb: Annotated[ 

514 bool, Arg(help="Log training metrics to Weights & Biases.") 

515 ] = False, 

516 ): 

517 """Train this YOLO object detector.""" 

518 return run_training_command( 

519 self, 

520 inputs, 

521 self._train, 

522 keep_log=keep_log, 

523 output=output, 

524 validation=validation, 

525 exclude=exclude, 

526 prepare_inputs=prepare_object_detection_inputs, 

527 include_unlabelled=True, 

528 model=model, 

529 epochs=epochs, 

530 batch=batch, 

531 imgsz=imgsz, 

532 early_stop=early_stop, 

533 preview=preview, 

534 workers=workers, 

535 pretrained=pretrained, 

536 optimizer=optimizer, 

537 seed=seed, 

538 deterministic=deterministic, 

539 single_cls=single_cls, 

540 rect=rect, 

541 cos_lr=cos_lr, 

542 close_mosaic=close_mosaic, 

543 amp=amp, 

544 fraction=fraction, 

545 freeze=freeze, 

546 multi_scale=multi_scale, 

547 lr0=lr0, 

548 lrf=lrf, 

549 momentum=momentum, 

550 weight_decay=weight_decay, 

551 warmup_epochs=warmup_epochs, 

552 warmup_momentum=warmup_momentum, 

553 warmup_bias_lr=warmup_bias_lr, 

554 box=box, 

555 cls=cls, 

556 cls_pw=cls_pw, 

557 dfl=dfl, 

558 nbs=nbs, 

559 hsv_h=hsv_h, 

560 hsv_s=hsv_s, 

561 hsv_v=hsv_v, 

562 degrees=degrees, 

563 translate=translate, 

564 scale=scale, 

565 shear=shear, 

566 perspective=perspective, 

567 flipud=flipud, 

568 fliplr=fliplr, 

569 bgr=bgr, 

570 mosaic=mosaic, 

571 mixup=mixup, 

572 cutmix=cutmix, 

573 time=time, 

574 save=save, 

575 save_period=save_period, 

576 cache=cache, 

577 device=device, 

578 verbose=verbose, 

579 resume=resume, 

580 compile=compile, 

581 distill_model=distill_model, 

582 dis=dis, 

583 plots=plots, 

584 wandb=wandb, 

585 ) 

586 

587 def _train( 

588 self, *images: ImageCrop, **kwargs 

589 ) -> DetMetrics | ClassifyMetrics | None: 

590 YOLO = get_yolo_class() 

591 

592 if kwargs.get("model"): 

593 self.model = kwargs["model"] 

594 self._init_model() 

595 log = kwargs.get("log", None) 

596 if log is None: 

597 raise ValueError( 

598 "Log directory must be specified in kwargs with key 'log'." 

599 ) 

600 log = kwargs["base_folder"] / log 

601 log.mkdir(parents=True, exist_ok=True) 

602 print(f"Training files will be saved to: {log}") 

603 classes = kwargs.pop("classes", []) 

604 

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

606 raise ValueError("Model must be a YOLO instance for training.") 

607 

608 print(f"Training with model: {self.model.model_name}") 

609 

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

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

612 train, val = self._prepare_assets(classes, log, validation, exclude, *images) 

613 return self._train_fold(train, val, classes, log, kwargs) 

614 

615 def _train_fold( 

616 self, 

617 train: list[str] | Path, 

618 val: list[str] | Path, 

619 classes: list[str], 

620 log: Path, 

621 kwargs: dict, 

622 fold_count: int | None = None, 

623 ) -> DetMetrics | ClassifyMetrics: 

624 YOLO = get_yolo_class() 

625 

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

627 raise ValueError("Expected model to be a YOLO instance for training.") 

628 kwargs["config_file"] = self._prepare_yolo_config( 

629 log, 

630 classes, 

631 train, 

632 val, 

633 ) 

634 kwargs["log"] = f"{log.name}" 

635 if fold_count: 

636 kwargs["log"] += f"_fold_{fold_count}" 

637 params = self._prepare_params(**kwargs) 

638 if kwargs.get("preview", False): 

639 self._preview_assets(log, classes) 

640 with ( 

641 change_dir(kwargs["base_folder"]), 

642 self._wandb_logging(kwargs.get("wandb", False)), 

643 ): 

644 results: DetMetrics | None = self.model.train(**params) 

645 if results is None: 

646 raise Exception("[red]Training failed, no results returned.[/red]") 

647 self._reload() 

648 return results 

649 

650 def _preview_assets(self, log: Path, classes: list[str]) -> None: 

651 try: 

652 from drawyolo.draw import draw_box_on_image_with_labels 

653 except ModuleNotFoundError as error: 

654 if error.name != "drawyolo": 

655 raise 

656 raise ImportError( 

657 'YOLO previews require drawyolo. Install with: ' 

658 'pip install "collectra[yolo]"' 

659 ) from error 

660 

661 preview_dir = log / "preview" 

662 preview_dir.mkdir(exist_ok=True) 

663 

664 for image_type in Image.image_types(): 

665 for image_file in log.glob(f"*{image_type}"): 

666 dest = preview_dir / image_file.name 

667 shutil.copy(image_file, dest) 

668 output_path = preview_dir / image_file.name 

669 draw_box_on_image_with_labels( 

670 image=image_file, 

671 output=output_path, 

672 labels=image_file.with_suffix(".txt"), 

673 classes=classes, 

674 ) 

675 

676 def _prepare_yolo_config( 

677 self, 

678 log: Path, 

679 classes: list[str], 

680 train: list[str] | Path, 

681 val: list[str] | Path, 

682 ) -> Path: 

683 

684 # Inspect exactly the prepared files referenced by the generated manifests. 

685 # The dataset, not the selected checkpoint, determines the model input size. 

686 train = [Path(path).name for path in train] 

687 val = [Path(path).name for path in val] 

688 detected = {} 

689 for name in dict.fromkeys([*train, *val]): 

690 path = log / name 

691 try: 

692 detected[name] = image_channel_count(path) 

693 except Exception as error: 

694 raise ValueError( 

695 f"Cannot determine input channels for {path}: {error}" 

696 ) from error 

697 if not detected: 

698 raise ValueError( 

699 "Cannot infer input channels: no prepared training or validation images" 

700 ) 

701 if len(set(detected.values())) != 1: 

702 details = ", ".join( 

703 f"{name}: {count} channels" for name, count in detected.items() 

704 ) 

705 raise ValueError( 

706 f"Inconsistent input channel counts in prepared dataset: {details}" 

707 ) 

708 channels = next(iter(detected.values())) 

709 

710 metadata = { 

711 "classes": classes, 

712 "train": train, 

713 "val": val, 

714 } 

715 

716 Path(log / "train.txt").write_text( 

717 "\n".join([f"./{Path(p).name}" for p in train]) 

718 ) 

719 Path(log / "val.txt").write_text("\n".join([f"./{Path(p).name}" for p in val])) 

720 

721 print( 

722 f"Saved {len(metadata['train'])} training images and {len(metadata['val'])} validation images to {log}" 

723 ) 

724 

725 self._check_distribution(log, metadata) 

726 

727 train_txt = "train.txt" 

728 val_txt = "val.txt" 

729 num_classes = len(metadata["classes"]) 

730 names = metadata["classes"] 

731 

732 config = ( 

733 f"train: {train_txt}\nval: {val_txt}\nnc: {num_classes}\nnames: {names}\n" 

734 f"channels: {channels}\n" 

735 ) 

736 config_file = log / "config.yml" 

737 Path(config_file).write_text(config) 

738 

739 return config_file 

740 

741 def _check_distribution( 

742 self, 

743 log: Path, 

744 metadata: dict | list[str], 

745 ) -> None: 

746 if not isinstance(metadata, dict): 

747 return 

748 

749 train_classes = {name: 0 for name in metadata["classes"]} 

750 val_classes = {name: 0 for name in metadata["classes"]} 

751 

752 for image_type in Image.image_types(): 

753 for image_file in log.glob(f"*{image_type}"): 

754 label_file = image_file.with_suffix(".txt") 

755 with open(label_file, "r") as f: 

756 label_lines = [ 

757 line.strip() for line in f.readlines() if line.strip() 

758 ] 

759 for line in label_lines: 

760 class_index = int(line.split()[0]) 

761 class_name = metadata["classes"][class_index] 

762 if image_file.name in metadata["train"]: 

763 train_classes[class_name] += 1 

764 elif image_file.name in metadata["val"]: 

765 val_classes[class_name] += 1 

766 

767 print_distribution_table( 

768 "Class Distribution", metadata["classes"], train_classes, val_classes 

769 ) 

770 

771 def _write_yolo_label(self, name_index: int, img: ImageCrop, f): 

772 if name_index == -1: 

773 f.write("") 

774 else: 

775 f.write( 

776 f"{name_index} {img.x_center:.6f} {img.y_center:.6f} {img.width_relative:.6f} {img.height_relative:.6f}\n" 

777 ) 

778 

779 def _prepare_assets( 

780 self, 

781 classes: list[str], 

782 log: Path, 

783 validation_flag: str, 

784 exclude_flag: str, 

785 *images: ImageCrop, 

786 ) -> tuple[list[str], list[str]] | tuple[Path, Path]: 

787 """Prepare assets for YOLO training. 

788 

789 Args: 

790 classes (list[str]): List of class names. 

791 log (Path): Path to store the assets for training. 

792 validation_flag (str): The value of the validation flag to identify validation images. 

793 exclude_flag (str): The value of the exclude flag to identify images to be excluded from training. 

794 images (ImageCrop): Variable number of ImageCrop instances to be prepared. 

795 

796 Returns: 

797 tuple[list[list[int]], list[str], list[str]]: A tuple containing the label matrix for iterative stratification, 

798 training image filenames, and validation image filenames. 

799 

800 """ 

801 train: list[str] = list() 

802 val: list[str] = list() 

803 for img in images: 

804 src = img.get_path() 

805 dst = log / src.name 

806 if isinstance(img, ImageCrop) and img.source_parent: 

807 dst = log / f"{src.stem}-{str(img.source_parent.id)}{src.suffix}" 

808 if not dst.exists(): 

809 if hasattr(img, "source_parent") and img.source_parent is not None: 

810 img.source_parent.save(dst) 

811 else: 

812 img.save(dst) 

813 if img.partition == validation_flag: 

814 val.append(dst.name) 

815 elif not exclude_flag or img.partition != exclude_flag: 

816 train.append(dst.name)