Coverage for collectra/tasks/object_detection/detr.py: 39.12%

1from __future__ import annotations

3from dataclasses import dataclass

4from pathlib import Path

5from typing import TYPE_CHECKING

7from collectra.cli import command

8from collectra.types.images import Image, ImageCrop

9from collectra.utils import change_dir

11from ...logger import get_logger

12from ..base import Task

13from collectra.utils import threading_locked

14from ..machine_learning.training import (

15 prepare_object_detection_inputs,

16 print_distribution_table,

17 run_training_command,

18)

20if TYPE_CHECKING:

21 from torch import nn

23logger = get_logger(__name__)

25__all__ = ["ObjectDetectionDETR", "DetectionTrainResult"]

28@dataclass

29class DetectionTrainResult:

30 save_dir: Path

31 results_dict: dict

34class ObjectDetectionDETR(Task):

35 """DETR task backed by Hugging Face Transformers."""

37 model: str | Path | nn.Module | None

38 original_model_path: str | Path | None = None

40 def __init__(

41 self,

42 name: str,

43 model: str | Path | nn.Module | None = None,

44 **kwargs,

45 ):

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

47 self._device = None

48 self._processor = None

49 self._categories: list[str] = []

50 self._model_name = "facebook/detr-resnet-50"

52 def _select_device(self):

53 import platform

55 import torch

57 return (

58 "mps"

59 if platform.system() == "Darwin"

60 else "cuda" if torch.cuda.is_available() else "cpu"

61 )

63 def _init_model(self) -> None:

64 from torch import nn

66 self._load()

67 if not isinstance(self.model, nn.Module):

68 raise ValueError("Model must be a torch.nn.Module instance")

70 def _reload(self) -> None:

71 if not self.original_model_path:

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

73 return

74 self.model = self.original_model_path

75 self._load()

77 def _load_model_from_hf(self, model_name: str) -> None:

78 from transformers import DetrForObjectDetection, DetrImageProcessor

80 self._model_name = model_name

81 self._processor = DetrImageProcessor.from_pretrained(model_name)

82 model = DetrForObjectDetection.from_pretrained(model_name)

83 self._categories = [

84 model.config.id2label[i] for i in sorted(model.config.id2label.keys())

85 ]

86 self.model = model.to(self._device)

87 self.model.eval()

89 def _load_model_from_checkpoint(self, checkpoint_path: Path) -> None:

90 import torch

91 from transformers import DetrConfig, DetrForObjectDetection, DetrImageProcessor

93 checkpoint = torch.load(checkpoint_path, map_location=self._device)

94 if isinstance(checkpoint, dict):

95 state_dict = checkpoint.get("model_state_dict", checkpoint)

96 classes = checkpoint.get("classes", [])

97 model_name = checkpoint.get("model_name", self._model_name)

98 else:

99 state_dict = checkpoint

100 classes = []

101 model_name = self._model_name

102

103 self._model_name = model_name

104 self._processor = DetrImageProcessor.from_pretrained(self._model_name)

105

106 if classes:

107 id2label = {idx: name for idx, name in enumerate(classes)}

108 label2id = {name: idx for idx, name in enumerate(classes)}

109 config = DetrConfig.from_pretrained(

110 self._model_name,

111 num_labels=len(classes),

112 id2label=id2label,

113 label2id=label2id,

114 )

115 model = DetrForObjectDetection.from_pretrained(

116 self._model_name,

117 config=config,

118 ignore_mismatched_sizes=True,

119 )

120 self._categories = classes

121 else:

122 model = DetrForObjectDetection.from_pretrained(self._model_name)

123 self._categories = [

124 model.config.id2label[i] for i in sorted(model.config.id2label.keys())

125 ]

126

127 model.load_state_dict(state_dict)

128 model.to(self._device)

129 model.eval()

130 self.model = model

131

132 def _load(self) -> None:

133 from torch import nn

134

135 if isinstance(self.model, nn.Module):

136 if self._device is None:

137 self._device = self._select_device()

138 self.model.to(self._device)

139 self.model.eval()

140 return

141

142 self._device = self._select_device()

143

144 if self.model is None or (

145 isinstance(self.model, str) and self.model == "default"

146 ):

147 self._load_model_from_hf("facebook/detr-resnet-50")

148 return

149

150 if not isinstance(self.model, (str, Path)):

151 raise ValueError(

152 "Model must be None, 'default', a path, a model id, or nn.Module"

153 )

154

155 model_value = str(self.model)

156 checkpoint_path = Path(model_value)

157 if checkpoint_path.exists() and checkpoint_path.is_file():

158 self.original_model_path = checkpoint_path

159 self._load_model_from_checkpoint(checkpoint_path)

160 return

161

162 # Treat as Hugging Face model id if local file path does not exist.

163 self._load_model_from_hf(model_value)

164

165 @threading_locked()

166 def run(self, *args: Image) -> list[Image]:

167 import torch

168

169 if len(args) != 1:

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

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

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

173

174 image = args[0]

175 self._init_model()

176 from torch import nn

177

178 if not isinstance(self.model, nn.Module) or self._processor is None:

179 raise ValueError("DETR model is not initialised properly.")

180

181 threshold = float(getattr(self, "threshold", 0.5))

182 pil_image = image.pil().convert("RGB")

183 width, height = pil_image.size

184

185 inputs = self._processor(images=pil_image, return_tensors="pt")

186 inputs = {k: v.to(self._device) for k, v in inputs.items()}

187

188 with torch.no_grad():

189 outputs = self.model(**inputs)

190

191 target_sizes = torch.tensor([[height, width]], device=self._device)

192 processed = self._processor.post_process_object_detection(

193 outputs,

194 threshold=threshold,

195 target_sizes=target_sizes,

196 )[0]

197

198 detections: list[Image] = []

199 for box, label in zip(processed["boxes"], processed["labels"]):

200 x1, y1, x2, y2 = box.detach().cpu().tolist()

201 x1 = max(0.0, min(float(width), float(x1)))

202 x2 = max(0.0, min(float(width), float(x2)))

203 y1 = max(0.0, min(float(height), float(y1)))

204 y2 = max(0.0, min(float(height), float(y2)))

205

206 bbox_width = x2 - x1

207 bbox_height = y2 - y1

208 if bbox_width <= 0 or bbox_height <= 0:

209 continue

210

211 x_center = (x1 + x2) / 2.0 / width

212 y_center = (y1 + y2) / 2.0 / height

213 width_relative = bbox_width / width

214 height_relative = bbox_height / height

215

216 class_idx = int(label.detach().cpu().item())

217 class_name = (

218 self._categories[class_idx]

219 if 0 <= class_idx < len(self._categories)

220 else str(class_idx)

221 )

222

223 detections.append(

224 image.make_crop(

225 x_center=x_center,

226 y_center=y_center,

227 width_relative=width_relative,

228 height_relative=height_relative,

229 orientation=image.orientation,

230 name=class_name,

231 )

232 )

233

234 print(f"Found {len(detections)} objects in the image.")

235 return detections

236

237 def _prepare_params(self, **kwargs) -> dict:

238 model_name = kwargs.get("model_name", self._model_name)

239 if isinstance(self.model, str) and self.model not in ("", "default"):

240 if not Path(self.model).exists():

241 model_name = self.model

242 return {

243 "name": kwargs["log"],

244 "project": kwargs.get("project", "runs/detr"),

245 "device": self._select_device(),

246 "epochs": int(kwargs.get("epochs", 1)),

247 "batch": int(kwargs.get("batch", 2)),

248 "workers": int(kwargs.get("workers", 0)),

249 "lr": float(kwargs.get("lr", 1e-4)),

250 "weight_decay": float(kwargs.get("weight_decay", 1e-4)),

251 "pretrained": bool(kwargs.get("pretrained", True)),

252 "model_name": model_name,

253 }

254

255 def _check_distribution(

256 self,

257 classes: list[str],

258 train_samples: list[dict],

259 val_samples: list[dict],

260 ) -> None:

261 train_classes = {name: 0 for name in classes}

262 val_classes = {name: 0 for name in classes}

263 label_to_name = {index: name for index, name in enumerate(classes)}

264

265 for sample in train_samples:

266 for label in sample["labels"]:

267 class_name = label_to_name.get(label, "")

268 if class_name:

269 train_classes[class_name] += 1

270

271 for sample in val_samples:

272 for label in sample["labels"]:

273 class_name = label_to_name.get(label, "")

274 if class_name:

275 val_classes[class_name] += 1

276

277 print_distribution_table(

278 "Class Distribution", classes, train_classes, val_classes

279 )

280

281 def _prepare_assets(

282 self,

283 classes: list[str],

284 validation_flag: str,

285 exclude_flag: str,

286 *images: ImageCrop,

287 ) -> tuple[list[dict], list[dict]]:

288 class_to_idx = {name: idx for idx, name in enumerate(classes)}

289 grouped: dict[Path, dict] = {}

290

291 for img in images:

292 if exclude_flag and img.partition == exclude_flag:

293 continue

294

295 base_img: Image | ImageCrop = img

296 if isinstance(img, ImageCrop):

297 if img.source_parent:

298 base_img = img.source_parent

299 if isinstance(img.source_parent, ImageCrop):

300 img.set_rel_to_src_parent()

301

302 image_path = base_img.get_path()

303 if image_path not in grouped:

304 grouped[image_path] = {

305 "image_path": image_path,

306 "partition": img.partition,

307 "boxes": [],

308 "labels": [],

309 }

310

311 if not isinstance(img, ImageCrop):

312 continue

313 if img.name not in class_to_idx:

314 continue

315

316 width, height = base_img.width, base_img.height

317 box_w = float(img.width_relative) * width

318 box_h = float(img.height_relative) * height

319 cx = float(img.x_center) * width

320 cy = float(img.y_center) * height

321

322 x1 = max(0.0, cx - box_w / 2.0)

323 y1 = max(0.0, cy - box_h / 2.0)

324 x2 = min(float(width), cx + box_w / 2.0)

325 y2 = min(float(height), cy + box_h / 2.0)

326 if x2 <= x1 or y2 <= y1:

327 continue

328

329 grouped[image_path]["boxes"].append([x1, y1, x2, y2])

330 grouped[image_path]["labels"].append(class_to_idx[img.name])

331

332 train_samples: list[dict] = []

333 val_samples: list[dict] = []

334 for sample in grouped.values():

335 if validation_flag and sample["partition"] == validation_flag:

336 val_samples.append(sample)

337 elif not exclude_flag or sample["partition"] != exclude_flag:

338 train_samples.append(sample)

339 return train_samples, val_samples

340

341 def _train_fold(

342 self,

343 train_samples: list[dict],

344 val_samples: list[dict],

345 classes: list[str],

346 log: Path,

347 kwargs: dict,

348 ) -> DetectionTrainResult:

349 import torch

350 from torch.utils.data import DataLoader, Dataset

351 from transformers import DetrConfig, DetrForObjectDetection, DetrImageProcessor

352

353 params_kwargs = dict(kwargs)

354 params_kwargs["log"] = f"{log.name}"

355 params = self._prepare_params(**params_kwargs)

356 self._device = params["device"]

357 self._model_name = params["model_name"]

358

359 self._processor = DetrImageProcessor.from_pretrained(self._model_name)

360

361 class DetrDataset(Dataset):

362 def __init__(self, samples: list[dict], processor: DetrImageProcessor):

363 self.samples = samples

364 self.processor = processor

365

366 def __len__(self):

367 return len(self.samples)

368

369 def __getitem__(self, idx: int):

370 sample = self.samples[idx]

371 from PIL import Image as PILImage

372

373 with PILImage.open(sample["image_path"]).convert("RGB") as pil_image:

374 annotations = {

375 "image_id": idx,

376 "annotations": [],

377 }

378 for box, label in zip(sample["boxes"], sample["labels"]):

379 x1, y1, x2, y2 = box

380 w = x2 - x1

381 h = y2 - y1

382 annotations["annotations"].append(

383 {

384 "bbox": [x1, y1, w, h],

385 "category_id": int(label),

386 "area": float(w * h),

387 "iscrowd": 0,

388 }

389 )

390

391 encoding = self.processor(

392 images=pil_image,

393 annotations=annotations,

394 return_tensors="pt",

395 )

396

397 pixel_values = encoding["pixel_values"].squeeze(0)

398 labels = encoding["labels"][0]

399 return pixel_values, labels

400

401 def collate_fn(batch):

402 import torch

403

404 pixel_values = [item[0] for item in batch]

405 labels = [item[1] for item in batch]

406

407 max_h = max(pv.shape[1] for pv in pixel_values)

408 max_w = max(pv.shape[2] for pv in pixel_values)

409 batch_size = len(pixel_values)

410

411 pixel_values_padded = torch.zeros(

412 (batch_size, 3, max_h, max_w),

413 dtype=pixel_values[0].dtype,

414 )

415 pixel_mask = torch.zeros((batch_size, max_h, max_w), dtype=torch.long)

416

417 for i, pv in enumerate(pixel_values):

418 h, w = pv.shape[1], pv.shape[2]

419 pixel_values_padded[i, :, :h, :w] = pv

420 pixel_mask[i, :h, :w] = 1

421

422 return {

423 "pixel_values": pixel_values_padded,

424 "pixel_mask": pixel_mask,

425 "labels": labels,

426 }

427

428 train_loader = DataLoader(

429 DetrDataset(train_samples, self._processor),

430 batch_size=params["batch"],

431 shuffle=True,

432 num_workers=params["workers"],

433 collate_fn=collate_fn,

434 )

435

436 val_loader = None

437 if len(val_samples) > 0:

438 val_loader = DataLoader(

439 DetrDataset(val_samples, self._processor),

440 batch_size=params["batch"],

441 shuffle=False,

442 num_workers=params["workers"],

443 collate_fn=collate_fn,

444 )

445

446 id2label = {idx: name for idx, name in enumerate(classes)}

447 label2id = {name: idx for idx, name in enumerate(classes)}

448

449 from torch import nn

450

451 if isinstance(self.model, nn.Module):

452 model = self.model

453 else:

454 config = DetrConfig.from_pretrained(

455 self._model_name,

456 num_labels=len(classes),

457 id2label=id2label,

458 label2id=label2id,

459 )

460 model = DetrForObjectDetection.from_pretrained(

461 self._model_name,

462 config=config,

463 ignore_mismatched_sizes=True,

464 )

465

466 model.to(self._device)

467 self.model = model

468 self._categories = classes

469

470 optimizer = torch.optim.AdamW(

471 model.parameters(),

472 lr=params["lr"],

473 weight_decay=params["weight_decay"],

474 )

475

476 best_score = float("inf")

477 best_metrics: dict = {}

478 weights_dir = log / "weights"

479 weights_dir.mkdir(parents=True, exist_ok=True)

480

481 wandb_run = None

482 if bool(kwargs.get("wandb", False)):

483 try:

484 import wandb

485

486 wandb_run = wandb.init(

487 project=params["project"],

488 name=params["name"],

489 dir=str(kwargs["base_folder"]),

490 config=params,

491 reinit=True,

492 )

493 except Exception as error:

494 logger.warning("W&B init skipped: %s", error)

495

496 def move_labels_to_device(batch_labels: list[dict]) -> list[dict]:

497 out = []

498 for label_dict in batch_labels:

499 out.append({k: v.to(self._device) for k, v in label_dict.items()})

500 return out

501

502 def run_epoch(data_loader: DataLoader, train_mode: bool) -> float:

503 if train_mode:

504 model.train()

505 else:

506 model.eval()

507 running_loss = 0.0

508 running_count = 0

509

510 for batch in data_loader:

511 pixel_values = batch["pixel_values"].to(self._device)

512 pixel_mask = batch["pixel_mask"].to(self._device)

513 labels = move_labels_to_device(batch["labels"])

514

515 if train_mode:

516 optimizer.zero_grad(set_to_none=True)

517

518 with torch.set_grad_enabled(train_mode):

519 outputs = model(

520 pixel_values=pixel_values,

521 pixel_mask=pixel_mask,

522 labels=labels,

523 )

524 loss = outputs.loss

525 if train_mode:

526 loss.backward()

527 optimizer.step()

528

529 running_loss += float(loss.detach().cpu().item())

530 running_count += 1

531

532 return running_loss / max(running_count, 1)

533

534 for epoch in range(params["epochs"]):

535 train_loss = run_epoch(train_loader, train_mode=True)

536 val_loss = run_epoch(val_loader, train_mode=False) if val_loader else None

537 metric_to_track = val_loss if val_loss is not None else train_loss

538

539 epoch_metrics = {"epoch": epoch + 1, "train/loss": train_loss}

540 if val_loss is not None:

541 epoch_metrics["val/loss"] = val_loss

542

543 if metric_to_track < best_score:

544 best_score = metric_to_track

545 best_metrics = epoch_metrics.copy()

546 torch.save(

547 {

548 "model_state_dict": model.state_dict(),

549 "classes": classes,

550 "epoch": epoch + 1,

551 "metrics": epoch_metrics,

552 "model_name": self._model_name,

553 },

554 weights_dir / "best.pt",

555 )

556

557 if wandb_run is not None:

558 wandb_run.log(epoch_metrics)

559

560 logger.info("DETR epoch %s metrics: %s", epoch + 1, epoch_metrics)

561

562 torch.save(

563 {

564 "model_state_dict": model.state_dict(),

565 "classes": classes,

566 "epoch": params["epochs"],

567 "metrics": best_metrics,

568 "model_name": self._model_name,

569 },

570 weights_dir / "last.pt",

571 )

572

573 if wandb_run is not None:

574 wandb_run.finish()

575

576 model.eval()

577 self.original_model_path = weights_dir / "best.pt"

578 self._reload()

579 return DetectionTrainResult(save_dir=log, results_dict=best_metrics)

580

581 @command

582 def train(

583 self,

584 inputs: list[str],

585 output: Path | None = None,

586 keep_log: bool = True,

587 validation: str = "",

588 exclude: str = "",

589 epochs: int = 1,

590 batch: int = 2,

591 workers: int = 0,

592 learning_rate: float = 1e-4,

593 weight_decay: float = 1e-4,

594 pretrained: bool = True,

595 model_name: str = "facebook/detr-resnet-50",

596 wandb: bool = False,

597 ):

598 """Train this DETR object detector."""

599 return run_training_command(

600 self,

601 inputs,

602 self._train,

603 keep_log=keep_log,

604 output=output,

605 validation=validation,

606 exclude=exclude,

607 prepare_inputs=prepare_object_detection_inputs,

608 include_unlabelled=True,

609 epochs=epochs,

610 batch=batch,

611 workers=workers,

612 lr=learning_rate,

613 weight_decay=weight_decay,

614 pretrained=pretrained,

615 model_name=model_name,

616 wandb=wandb,

617 )

618

619 def _train(self, *images: ImageCrop, **kwargs) -> DetectionTrainResult:

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

621 if log is None:

622 raise ValueError(

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

624 )

625 if "base_folder" not in kwargs:

626 raise ValueError(

627 "Base folder must be specified in kwargs with key 'base_folder'."

628 )

629

630 log = kwargs["base_folder"] / log

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

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

633

634 classes = kwargs.get("classes", [])

635 if not classes:

636 classes = sorted({img.name for img in images if isinstance(img, ImageCrop)})

637 if not classes:

638 raise ValueError("No classes provided/found for DETR training.")

639

640 print("Training with model: DETR")

641

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

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

644 train_samples, val_samples = self._prepare_assets(

645 classes,

646 validation,

647 exclude,

648 *images,

649 )

650 self._check_distribution(classes, train_samples, val_samples)

651

652 with change_dir(kwargs["base_folder"]):

653 results = self._train_fold(train_samples, val_samples, classes, log, kwargs)

654

655 print(f"Best DETR model saved to: {results.save_dir / 'weights' / 'best.pt'}")

656 return results