From 3cc93ebdb767e29e1542a18389f5a89284e2fbba Mon Sep 17 00:00:00 2001 From: NikhilReddy Date: Tue, 9 Jun 2026 15:11:20 +0530 Subject: [PATCH] Add DINOv3 OWEED ViTDet training pipeline Add a frozen Hugging Face DINOv3 ViT-B/16 backbone for ViTDet-style Mask R-CNN training while keeping the detector heads and SimpleFeaturePyramid architecture compatible with Detectron2 LazyConfig. Add OWEED-specific training and evaluation configs for the original, tiled v0, and tiled v1 datasets, including bbox-only evaluation for memory-sensitive validation, best-checkpoint tracking by bbox AP, and a mask-eval config for full COCO bbox+segm metrics. Add a COCO tiling utility for offline 1024px tiled datasets with overlap support, plus a full-validation inference script that can dump annotated images, log timing/GPU memory, and optionally export low-memory full-resolution contours scaled from compact ROI masks instead of materializing dense 4K masks on GPU. --- detectron2/modeling/backbone/__init__.py | 1 + detectron2/modeling/backbone/hf_dinov3.py | 108 +++++ detectron2/modeling/meta_arch/__init__.py | 1 + .../modeling/meta_arch/bbox_only_eval_rcnn.py | 37 ++ .../COCO/mask_rcnn_dinov3_vitb16_100ep.py | 21 + .../OWEED/mask_rcnn_dinov3_vitb16_frozen.py | 64 +++ ...cnn_dinov3_vitb16_frozen_tiled_keepfrag.py | 99 ++++ ..._vitb16_frozen_tiled_v1_from_v0_best_ap.py | 102 ++++ ...itb16_frozen_tiled_v1_mask_eval_best_ap.py | 28 ++ projects/ViTDet/tools/create_tiled_coco.py | 323 +++++++++++++ tools/lazyconfig_train_net.py | 9 + tools/oweed_full_val_infer_annotate.py | 441 ++++++++++++++++++ 12 files changed, 1234 insertions(+) create mode 100644 detectron2/modeling/backbone/hf_dinov3.py create mode 100644 detectron2/modeling/meta_arch/bbox_only_eval_rcnn.py create mode 100644 projects/ViTDet/configs/COCO/mask_rcnn_dinov3_vitb16_100ep.py create mode 100644 projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen.py create mode 100644 projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_keepfrag.py create mode 100644 projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_v1_from_v0_best_ap.py create mode 100644 projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_v1_mask_eval_best_ap.py create mode 100644 projects/ViTDet/tools/create_tiled_coco.py create mode 100644 tools/oweed_full_val_infer_annotate.py diff --git a/detectron2/modeling/backbone/__init__.py b/detectron2/modeling/backbone/__init__.py index 5b3358a406..340cddde9d 100644 --- a/detectron2/modeling/backbone/__init__.py +++ b/detectron2/modeling/backbone/__init__.py @@ -13,6 +13,7 @@ BottleneckBlock, ) from .vit import ViT, SimpleFeaturePyramid, get_vit_lr_decay_rate +from .hf_dinov3 import HFDINOv3ViT from .mvit import MViT from .swin import SwinTransformer diff --git a/detectron2/modeling/backbone/hf_dinov3.py b/detectron2/modeling/backbone/hf_dinov3.py new file mode 100644 index 0000000000..189f21f090 --- /dev/null +++ b/detectron2/modeling/backbone/hf_dinov3.py @@ -0,0 +1,108 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +from typing import Any, Dict, Optional + +import torch + +from detectron2.layers import ShapeSpec + +from .backbone import Backbone + +__all__ = ["HFDINOv3ViT"] + + +class HFDINOv3ViT(Backbone): + """ + Hugging Face DINOv3 ViT backbone adapter for ViTDet-style models. + + The adapter exposes the DINOv3 patch tokens as a single spatial feature map + with stride equal to the patch size, matching the contract expected by + :class:`SimpleFeaturePyramid`. + """ + + def __init__( + self, + model_name: str = "facebook/dinov3-vitb16-pretrain-lvd1689m", + *, + out_feature: str = "last_feat", + pretrained: bool = True, + freeze: bool = True, + config_kwargs: Optional[Dict[str, Any]] = None, + model_kwargs: Optional[Dict[str, Any]] = None, + ): + super().__init__() + config_kwargs = {} if config_kwargs is None else dict(config_kwargs) + model_kwargs = {} if model_kwargs is None else dict(model_kwargs) + + try: + from transformers import DINOv3ViTConfig, DINOv3ViTModel + except ImportError as exc: + raise ImportError( + "HFDINOv3ViT requires Hugging Face Transformers. " + "Install it with `pip install transformers`." + ) from exc + + if pretrained: + self.model = DINOv3ViTModel.from_pretrained(model_name, **model_kwargs) + else: + self.model = DINOv3ViTModel(DINOv3ViTConfig(**config_kwargs)) + + self.out_feature = out_feature + self.freeze_encoder = freeze + + cfg = self.model.config + self.patch_size = int(cfg.patch_size) + self.num_prefix_tokens = 1 + int(getattr(cfg, "num_register_tokens", 0)) + self._out_features = [out_feature] + self._out_feature_channels = {out_feature: int(cfg.hidden_size)} + self._out_feature_strides = {out_feature: self.patch_size} + + if self.freeze_encoder: + for parameter in self.model.parameters(): + parameter.requires_grad = False + self.model.eval() + + def train(self, mode: bool = True): + super().train(mode) + if self.freeze_encoder: + self.model.eval() + return self + + def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: + height, width = x.shape[-2:] + if height % self.patch_size != 0 or width % self.patch_size != 0: + raise ValueError( + f"Input height/width must be divisible by patch_size={self.patch_size}; " + f"got {(height, width)}." + ) + + outputs = self.model(pixel_values=x, return_dict=True) + hidden = outputs.last_hidden_state + + if hidden.dim() == 4: + features = hidden + else: + grid_h, grid_w = height // self.patch_size, width // self.patch_size + expected_tokens = grid_h * grid_w + if hidden.shape[1] == expected_tokens: + patch_tokens = hidden + elif hidden.shape[1] >= self.num_prefix_tokens + expected_tokens: + patch_tokens = hidden[ + :, self.num_prefix_tokens : self.num_prefix_tokens + expected_tokens, : + ] + else: + raise ValueError( + "DINOv3 output token count cannot be reshaped to a spatial feature map: " + f"got {hidden.shape[1]} tokens, expected {expected_tokens} patch tokens " + f"after {self.num_prefix_tokens} prefix tokens." + ) + features = patch_tokens.reshape(x.shape[0], grid_h, grid_w, -1).permute(0, 3, 1, 2) + + return {self.out_feature: features} + + def output_shape(self): + return { + self.out_feature: ShapeSpec( + channels=self._out_feature_channels[self.out_feature], + stride=self._out_feature_strides[self.out_feature], + ) + } diff --git a/detectron2/modeling/meta_arch/__init__.py b/detectron2/modeling/meta_arch/__init__.py index 6b06681570..5bd3b55c3b 100644 --- a/detectron2/modeling/meta_arch/__init__.py +++ b/detectron2/modeling/meta_arch/__init__.py @@ -7,6 +7,7 @@ # import all the meta_arch, so they will be registered from .rcnn import GeneralizedRCNN, ProposalNetwork +from .bbox_only_eval_rcnn import GeneralizedRCNNWithBBoxOnlyEval from .dense_detector import DenseDetector from .retinanet import RetinaNet from .fcos import FCOS diff --git a/detectron2/modeling/meta_arch/bbox_only_eval_rcnn.py b/detectron2/modeling/meta_arch/bbox_only_eval_rcnn.py new file mode 100644 index 0000000000..0221210135 --- /dev/null +++ b/detectron2/modeling/meta_arch/bbox_only_eval_rcnn.py @@ -0,0 +1,37 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +from typing import Dict, List, Optional + +import torch + +from detectron2.structures import Instances + +from .rcnn import GeneralizedRCNN + +__all__ = ["GeneralizedRCNNWithBBoxOnlyEval"] + + +class GeneralizedRCNNWithBBoxOnlyEval(GeneralizedRCNN): + """ + GeneralizedRCNN variant that trains all configured heads, but skips mask + prediction during inference/evaluation. + + This is useful for bbox-only COCO evaluation of mask models on very large + images, where postprocessing predicted masks would materialize full-size + bitmasks for every detection. + """ + + def inference( + self, + batched_inputs: List[Dict[str, torch.Tensor]], + detected_instances: Optional[List[Instances]] = None, + do_postprocess: bool = True, + ): + old_mask_on = getattr(self.roi_heads, "mask_on", None) + if old_mask_on is None: + return super().inference(batched_inputs, detected_instances, do_postprocess) + + self.roi_heads.mask_on = False + try: + return super().inference(batched_inputs, detected_instances, do_postprocess) + finally: + self.roi_heads.mask_on = old_mask_on diff --git a/projects/ViTDet/configs/COCO/mask_rcnn_dinov3_vitb16_100ep.py b/projects/ViTDet/configs/COCO/mask_rcnn_dinov3_vitb16_100ep.py new file mode 100644 index 0000000000..d1074470f0 --- /dev/null +++ b/projects/ViTDet/configs/COCO/mask_rcnn_dinov3_vitb16_100ep.py @@ -0,0 +1,21 @@ +from detectron2.config import LazyCall as L +from detectron2.modeling.backbone import HFDINOv3ViT + +from .mask_rcnn_vitdet_b_100ep import dataloader, lr_multiplier, model, optimizer, train + + +# Keep the ViTDet Mask R-CNN architecture unchanged, but replace the MAE ViT-B +# encoder with a frozen Hugging Face DINOv3 ViT-B/16 encoder. +model.backbone.net = L(HFDINOv3ViT)( + model_name="facebook/dinov3-vitb16-pretrain-lvd1689m", + out_feature="last_feat", + pretrained=True, + freeze=True, +) + +# Initialize trainable detector components from the COCO ViTDet-B detector +# checkpoint. The DINOv3 encoder loads its own Hugging Face weights above. +train.init_checkpoint = ( + "https://dl.fbaipublicfiles.com/detectron2/ViTDet/COCO/" + "mask_rcnn_vitdet_b/f325346929/model_final_61ccd1.pkl" +) diff --git a/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen.py b/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen.py new file mode 100644 index 0000000000..bbd91aec90 --- /dev/null +++ b/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen.py @@ -0,0 +1,64 @@ +from detectron2.data.datasets import register_coco_instances +from detectron2.modeling.meta_arch import GeneralizedRCNNWithBBoxOnlyEval + +from ..COCO.mask_rcnn_dinov3_vitb16_100ep import ( + dataloader, + lr_multiplier, + model, + optimizer, + train, +) + + +_DATA_ROOT = "/home/nikhileswara/Datasets/oweed_structure_v0_100_images_coco" +_TRAIN_NAME = "oweed_structure_v0_train" +_VAL_NAME = "oweed_structure_v0_val" + +register_coco_instances( + _TRAIN_NAME, + {}, + f"{_DATA_ROOT}/annotations/instances_train2017.json", + f"{_DATA_ROOT}/train2017", +) +register_coco_instances( + _VAL_NAME, + {}, + f"{_DATA_ROOT}/annotations/instances_val2017.json", + f"{_DATA_ROOT}/val2017", +) + +dataloader.train.dataset.names = _TRAIN_NAME +# 2 RTX 4090s: total_batch_size is global across DDP workers, so this gives +# two 1024x1024 LSJ images per GPU when launched with --num-gpus 2. +dataloader.train.total_batch_size = 4 +dataloader.train.num_workers = 8 +dataloader.train.pin_memory = True +dataloader.train.persistent_workers = True +dataloader.train.prefetch_factor = 4 +dataloader.test.dataset.names = _VAL_NAME +dataloader.test.num_workers = 8 + +dataloader.evaluator.dataset_name = _VAL_NAME +dataloader.evaluator.tasks = ("bbox",) + +model._target_ = GeneralizedRCNNWithBBoxOnlyEval +model.roi_heads.num_classes = 44 +model.roi_heads.box_predictor.test_score_thresh = 0.02 +model.roi_heads.box_predictor.test_topk_per_image = 500 + +# The dataset has up to ~460 annotated instances per image, so keep more +# proposals and final detections than the COCO defaults. +model.proposal_generator.pre_nms_topk = (4000, 4000) +model.proposal_generator.post_nms_topk = (2000, 2000) + +optimizer.lr = 1e-4 + +train.output_dir = "./output/oweed_dinov3_vitb16_frozen_mask_rcnn" +train.max_iter = 5000 +train.eval_period = 1000 +train.checkpointer.period = 500 +train.checkpointer.max_to_keep = 10 + +lr_multiplier.scheduler.milestones = [3500, 4500] +lr_multiplier.scheduler.num_updates = train.max_iter +lr_multiplier.warmup_length = 100 / train.max_iter diff --git a/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_keepfrag.py b/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_keepfrag.py new file mode 100644 index 0000000000..e9e6c16410 --- /dev/null +++ b/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_keepfrag.py @@ -0,0 +1,99 @@ +from fvcore.common.param_scheduler import MultiStepParamScheduler + +import detectron2.data.transforms as T +from detectron2.config import LazyCall as L +from detectron2.data.datasets import register_coco_instances +from detectron2.modeling.meta_arch import GeneralizedRCNNWithBBoxOnlyEval +from detectron2.solver import WarmupParamScheduler + +from ..COCO.mask_rcnn_dinov3_vitb16_100ep import ( + dataloader, + lr_multiplier, + model, + optimizer, + train, +) + + +_DATA_ROOT = "/home/nikhileswara/Datasets/oweed_structure_v0_100_images_coco_tiled_1024_s512_keepfrag" +_TRAIN_NAME = "oweed_structure_v0_tiled_keepfrag_train" +_VAL_NAME = "oweed_structure_v0_tiled_keepfrag_val" + +register_coco_instances( + _TRAIN_NAME, + {}, + f"{_DATA_ROOT}/annotations/instances_train2017.json", + f"{_DATA_ROOT}/train2017", +) +register_coco_instances( + _VAL_NAME, + {}, + f"{_DATA_ROOT}/annotations/instances_val2017.json", + f"{_DATA_ROOT}/val2017", +) + +dataloader.train.dataset.names = _TRAIN_NAME +dataloader.train.total_batch_size = 8 +dataloader.train.num_workers = 8 +dataloader.train.pin_memory = True +dataloader.train.persistent_workers = True +dataloader.train.prefetch_factor = 4 +dataloader.train.mapper.augmentations = [ + L(T.RandomFlip)(horizontal=True), + L(T.ResizeScale)( + min_scale=0.8, + max_scale=1.25, + target_height=1024, + target_width=1024, + ), + L(T.FixedSizeCrop)(crop_size=(1024, 1024), pad=True), +] +dataloader.train.mapper.image_format = "RGB" +dataloader.train.mapper.recompute_boxes = True + +dataloader.test.dataset.names = _VAL_NAME +dataloader.test.num_workers = 8 +dataloader.test.mapper.augmentations = [ + L(T.ResizeShortestEdge)(short_edge_length=1024, max_size=1024), +] + +dataloader.evaluator.dataset_name = _VAL_NAME +dataloader.evaluator.tasks = ("bbox",) +dataloader.evaluator.max_dets_per_image = 500 + +model._target_ = GeneralizedRCNNWithBBoxOnlyEval +model.roi_heads.num_classes = 44 +model.roi_heads.batch_size_per_image = 1024 +model.roi_heads.positive_fraction = 0.5 +model.roi_heads.box_predictor.test_score_thresh = 0.02 +model.roi_heads.box_predictor.test_nms_thresh = 0.7 +model.roi_heads.box_predictor.test_topk_per_image = 500 + +model.proposal_generator.anchor_generator.sizes = [[8], [16], [32], [64], [128]] +model.proposal_generator.batch_size_per_image = 512 +model.proposal_generator.positive_fraction = 0.5 +model.proposal_generator.pre_nms_topk = (4000, 4000) +model.proposal_generator.post_nms_topk = (2000, 2000) + +optimizer.lr = 1e-4 + +train.output_dir = "./output/oweed_tiled_keepfrag_dinov3_vitb16_frozen_mask_rcnn" +train.max_iter = 50000 +train.eval_period = 2000 +train.checkpointer.period = 2000 +train.checkpointer.max_to_keep = 10 +train.best_checkpointer = dict( + val_metric="bbox/APs", + mode="max", + file_prefix="model_best_bbox_aps", +) + +lr_multiplier = L(WarmupParamScheduler)( + scheduler=L(MultiStepParamScheduler)( + values=[1.0, 0.1, 0.01], + milestones=[35000, 45000], + num_updates=train.max_iter, + ), + warmup_length=500 / train.max_iter, + warmup_factor=0.001, +) diff --git a/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_v1_from_v0_best_ap.py b/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_v1_from_v0_best_ap.py new file mode 100644 index 0000000000..e79bfcd5fa --- /dev/null +++ b/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_v1_from_v0_best_ap.py @@ -0,0 +1,102 @@ +from fvcore.common.param_scheduler import MultiStepParamScheduler + +import detectron2.data.transforms as T +from detectron2.config import LazyCall as L +from detectron2.data.datasets import register_coco_instances +from detectron2.modeling.meta_arch import GeneralizedRCNNWithBBoxOnlyEval +from detectron2.solver import WarmupParamScheduler + +from ..COCO.mask_rcnn_dinov3_vitb16_100ep import ( + dataloader, + lr_multiplier, + model, + optimizer, + train, +) + + +_DATA_ROOT = "/home/nikhileswara/Datasets/oweed_structure_v1_177_images_coco_tiled" +_TRAIN_NAME = "oweed_structure_v1_tiled_train" +_VAL_NAME = "oweed_structure_v1_tiled_val" + +register_coco_instances( + _TRAIN_NAME, + {}, + f"{_DATA_ROOT}/annotations/instances_train2017.json", + f"{_DATA_ROOT}/train2017", +) +register_coco_instances( + _VAL_NAME, + {}, + f"{_DATA_ROOT}/annotations/instances_val2017.json", + f"{_DATA_ROOT}/val2017", +) + +dataloader.train.dataset.names = _TRAIN_NAME +dataloader.train.total_batch_size = 16 +dataloader.train.num_workers = 16 +dataloader.train.pin_memory = True +dataloader.train.persistent_workers = True +dataloader.train.prefetch_factor = 4 +dataloader.train.mapper.augmentations = [ + L(T.RandomFlip)(horizontal=True), + L(T.ResizeScale)( + min_scale=0.8, + max_scale=1.25, + target_height=1024, + target_width=1024, + ), + L(T.FixedSizeCrop)(crop_size=(1024, 1024), pad=True), +] +dataloader.train.mapper.image_format = "RGB" +dataloader.train.mapper.recompute_boxes = True + +dataloader.test.dataset.names = _VAL_NAME +dataloader.test.num_workers = 16 +dataloader.test.mapper.augmentations = [ + L(T.ResizeShortestEdge)(short_edge_length=1024, max_size=1024), +] + +dataloader.evaluator.dataset_name = _VAL_NAME +dataloader.evaluator.tasks = ("bbox",) +dataloader.evaluator.max_dets_per_image = 500 + +model._target_ = GeneralizedRCNNWithBBoxOnlyEval +model.roi_heads.num_classes = 44 +model.roi_heads.batch_size_per_image = 1024 +model.roi_heads.positive_fraction = 0.5 +model.roi_heads.box_predictor.test_score_thresh = 0.02 +model.roi_heads.box_predictor.test_nms_thresh = 0.7 +model.roi_heads.box_predictor.test_topk_per_image = 500 + +model.proposal_generator.anchor_generator.sizes = [[8], [16], [32], [64], [128]] +model.proposal_generator.batch_size_per_image = 512 +model.proposal_generator.positive_fraction = 0.5 +model.proposal_generator.pre_nms_topk = (4000, 4000) +model.proposal_generator.post_nms_topk = (2000, 2000) + +optimizer.lr = 5e-5 + +train.init_checkpoint = ( + "./output/oweed_tiled_keepfrag_dinov3_vitb16_frozen_mask_rcnn/model_0047999.pth" +) +train.output_dir = "./output/oweed_v1_tiled_from_v0_best_ap_dinov3_vitb16_frozen_mask_rcnn" +train.max_iter = 15000 +train.eval_period = 1000 +train.checkpointer.period = 1000 +train.checkpointer.max_to_keep = 10 +train.best_checkpointer = dict( + val_metric="bbox/AP", + mode="max", + file_prefix="model_best_bbox_ap", +) + +lr_multiplier = L(WarmupParamScheduler)( + scheduler=L(MultiStepParamScheduler)( + values=[1.0, 0.1, 0.01], + milestones=[10000, 13500], + num_updates=train.max_iter, + ), + warmup_length=300 / train.max_iter, + warmup_factor=0.001, +) diff --git a/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_v1_mask_eval_best_ap.py b/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_v1_mask_eval_best_ap.py new file mode 100644 index 0000000000..f73cbfc9b2 --- /dev/null +++ b/projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_v1_mask_eval_best_ap.py @@ -0,0 +1,28 @@ +from detectron2.modeling.meta_arch import GeneralizedRCNN + +from .mask_rcnn_dinov3_vitb16_frozen_tiled_v1_from_v0_best_ap import ( + dataloader, + lr_multiplier, + model, + optimizer, + train, +) + + +_EVAL_OUTPUT_DIR = ( + "./output/oweed_v1_tiled_from_v0_best_ap_dinov3_vitb16_frozen_mask_rcnn/" + "eval_mask_best_bbox_ap" +) + +# Use the normal Mask R-CNN inference path for this one-off evaluation so COCO +# metrics include both bounding boxes and instance masks. +model._target_ = GeneralizedRCNN + +dataloader.evaluator.tasks = ("bbox", "segm") +dataloader.evaluator.output_dir = _EVAL_OUTPUT_DIR + +train.init_checkpoint = ( + "./output/oweed_v1_tiled_from_v0_best_ap_dinov3_vitb16_frozen_mask_rcnn/" + "model_best_bbox_ap.pth" +) +train.output_dir = _EVAL_OUTPUT_DIR diff --git a/projects/ViTDet/tools/create_tiled_coco.py b/projects/ViTDet/tools/create_tiled_coco.py new file mode 100644 index 0000000000..7cd6d96927 --- /dev/null +++ b/projects/ViTDet/tools/create_tiled_coco.py @@ -0,0 +1,323 @@ +#!/usr/bin/env python3 +"""Create an offline tiled COCO instance segmentation dataset. + +This is intended for high-resolution dense-object datasets where resizing the +full image makes small instances too tiny for the detector. It preserves the +existing train/val split and clips polygon annotations into each tile. +""" + +from __future__ import annotations + +import argparse +import json +from collections import defaultdict +from pathlib import Path +from typing import Iterable + +from PIL import Image +from shapely.geometry import GeometryCollection, MultiPolygon, Polygon, box + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--src-root", + default="/home/nikhileswara/Datasets/oweed_structure_v0_100_images_coco", + help="Source COCO root containing train2017, val2017 and annotations/.", + ) + parser.add_argument( + "--out-root", + default="/home/nikhileswara/Datasets/oweed_structure_v0_100_images_coco_tiled_1024_s512", + help="Output tiled COCO root.", + ) + parser.add_argument("--tile-size", type=int, default=1024) + parser.add_argument("--stride", type=int, default=512) + parser.add_argument( + "--min-area", + type=float, + default=16.0, + help="Discard clipped annotation fragments smaller than this tile-local area.", + ) + parser.add_argument( + "--min-visible-ratio", + type=float, + default=0.05, + help="Discard clipped fragments with less than this fraction of original area.", + ) + parser.add_argument( + "--jpeg-quality", + type=int, + default=95, + help="JPEG quality for saved tiles.", + ) + parser.add_argument( + "--include-empty-train-tiles", + action="store_true", + help="Keep train tiles with no annotations. Val tiles are always kept.", + ) + return parser.parse_args() + + +def tile_starts(length: int, tile_size: int, stride: int) -> list[int]: + if length <= tile_size: + return [0] + starts = list(range(0, length - tile_size + 1, stride)) + last = length - tile_size + if starts[-1] != last: + starts.append(last) + return starts + + +def iter_polygons(geom) -> Iterable[Polygon]: + if geom.is_empty: + return + if isinstance(geom, Polygon): + yield geom + elif isinstance(geom, (MultiPolygon, GeometryCollection)): + for part in geom.geoms: + yield from iter_polygons(part) + + +def make_valid_polygon(points: list[tuple[float, float]]): + if len(points) < 3: + return None + geom = Polygon(points) + if geom.is_empty: + return None + if not geom.is_valid: + geom = geom.buffer(0) + if geom.is_empty: + return None + return geom + + +def segmentation_to_geometry(segmentation): + if not isinstance(segmentation, list): + raise ValueError("Only polygon COCO segmentations are supported by this tiler.") + + polygons = [] + for poly in segmentation: + if len(poly) < 6: + continue + points = [(float(poly[i]), float(poly[i + 1])) for i in range(0, len(poly), 2)] + geom = make_valid_polygon(points) + if geom is not None: + polygons.append(geom) + + if not polygons: + return GeometryCollection() + if len(polygons) == 1: + return polygons[0] + return MultiPolygon(polygons).buffer(0) + + +def polygon_to_coco(poly: Polygon, tile_x: int, tile_y: int, tile_size: int) -> list[float] | None: + coords = [] + for x, y in list(poly.exterior.coords)[:-1]: + lx = min(max(float(x) - tile_x, 0.0), float(tile_size)) + ly = min(max(float(y) - tile_y, 0.0), float(tile_size)) + coords.extend([lx, ly]) + + if len(coords) < 6: + return None + unique_points = {(round(coords[i], 3), round(coords[i + 1], 3)) for i in range(0, len(coords), 2)} + if len(unique_points) < 3: + return None + return coords + + +def bbox_intersects_tile(bbox_xywh, tile_x: int, tile_y: int, tile_size: int) -> bool: + x, y, w, h = [float(v) for v in bbox_xywh] + return not ( + x + w <= tile_x + or y + h <= tile_y + or x >= tile_x + tile_size + or y >= tile_y + tile_size + ) + + +def process_split( + split: str, + src_root: Path, + out_root: Path, + tile_size: int, + stride: int, + min_area: float, + min_visible_ratio: float, + jpeg_quality: int, + include_empty_train_tiles: bool, +) -> dict: + ann_path = src_root / "annotations" / f"instances_{split}2017.json" + image_dir = src_root / f"{split}2017" + out_image_dir = out_root / f"{split}2017" + out_ann_path = out_root / "annotations" / f"instances_{split}2017.json" + + out_image_dir.mkdir(parents=True, exist_ok=True) + out_ann_path.parent.mkdir(parents=True, exist_ok=True) + + with ann_path.open() as f: + coco = json.load(f) + + anns_by_image = defaultdict(list) + for ann in coco["annotations"]: + anns_by_image[ann["image_id"]].append(ann) + + out_images = [] + out_annotations = [] + next_image_id = 1 + next_ann_id = 1 + generated_tiles = 0 + kept_tiles = 0 + skipped_empty_train_tiles = 0 + skipped_fragments = 0 + + for image_info in coco["images"]: + image_path = image_dir / image_info["file_name"] + with Image.open(image_path) as img: + img = img.convert("RGB") + width, height = img.size + x_starts = tile_starts(width, tile_size, stride) + y_starts = tile_starts(height, tile_size, stride) + + image_annotations = anns_by_image[image_info["id"]] + geom_cache = {} + + for tile_y in y_starts: + for tile_x in x_starts: + generated_tiles += 1 + tile_geom = box(tile_x, tile_y, tile_x + tile_size, tile_y + tile_size) + tile_annotations = [] + + for ann in image_annotations: + if not bbox_intersects_tile(ann["bbox"], tile_x, tile_y, tile_size): + continue + + if ann["id"] not in geom_cache: + geom_cache[ann["id"]] = segmentation_to_geometry(ann["segmentation"]) + + clipped = geom_cache[ann["id"]].intersection(tile_geom) + clipped_area = float(clipped.area) + original_area = max(float(ann.get("area", clipped_area)), 1e-6) + if clipped_area < min_area or clipped_area / original_area < min_visible_ratio: + skipped_fragments += 1 + continue + + segmentations = [] + for poly in iter_polygons(clipped): + if float(poly.area) < min_area: + continue + coco_poly = polygon_to_coco(poly, tile_x, tile_y, tile_size) + if coco_poly is not None: + segmentations.append(coco_poly) + + if not segmentations: + skipped_fragments += 1 + continue + + minx, miny, maxx, maxy = clipped.bounds + minx = min(max(float(minx) - tile_x, 0.0), float(tile_size)) + miny = min(max(float(miny) - tile_y, 0.0), float(tile_size)) + maxx = min(max(float(maxx) - tile_x, 0.0), float(tile_size)) + maxy = min(max(float(maxy) - tile_y, 0.0), float(tile_size)) + if maxx <= minx or maxy <= miny: + skipped_fragments += 1 + continue + + tile_annotations.append( + { + "id": next_ann_id, + "image_id": next_image_id, + "category_id": ann["category_id"], + "segmentation": segmentations, + "bbox": [minx, miny, maxx - minx, maxy - miny], + "area": clipped_area, + "iscrowd": ann.get("iscrowd", 0), + "original_annotation_id": ann["id"], + "original_image_id": image_info["id"], + "tile_x": tile_x, + "tile_y": tile_y, + } + ) + next_ann_id += 1 + + if split == "train" and not include_empty_train_tiles and not tile_annotations: + skipped_empty_train_tiles += 1 + continue + + stem = Path(image_info["file_name"]).stem + tile_file_name = f"{stem}_x{tile_x:04d}_y{tile_y:04d}.jpg" + tile = img.crop((tile_x, tile_y, tile_x + tile_size, tile_y + tile_size)) + tile.save(out_image_dir / tile_file_name, quality=jpeg_quality) + + out_images.append( + { + "id": next_image_id, + "file_name": tile_file_name, + "width": tile_size, + "height": tile_size, + "original_image_id": image_info["id"], + "original_file_name": image_info["file_name"], + "tile_x": tile_x, + "tile_y": tile_y, + "tile_size": tile_size, + "stride": stride, + } + ) + out_annotations.extend(tile_annotations) + next_image_id += 1 + kept_tiles += 1 + + out_coco = { + "images": out_images, + "annotations": out_annotations, + "categories": coco["categories"], + "tiled_dataset": { + "source_root": str(src_root), + "split": split, + "tile_size": tile_size, + "stride": stride, + "min_area": min_area, + "min_visible_ratio": min_visible_ratio, + }, + } + + with out_ann_path.open("w") as f: + json.dump(out_coco, f) + + return { + "split": split, + "generated_tiles": generated_tiles, + "kept_tiles": kept_tiles, + "skipped_empty_train_tiles": skipped_empty_train_tiles, + "annotations": len(out_annotations), + "skipped_fragments": skipped_fragments, + "annotation_path": str(out_ann_path), + } + + +def main() -> None: + args = parse_args() + src_root = Path(args.src_root) + out_root = Path(args.out_root) + + summaries = [] + for split in ("train", "val"): + summaries.append( + process_split( + split=split, + src_root=src_root, + out_root=out_root, + tile_size=args.tile_size, + stride=args.stride, + min_area=args.min_area, + min_visible_ratio=args.min_visible_ratio, + jpeg_quality=args.jpeg_quality, + include_empty_train_tiles=args.include_empty_train_tiles, + ) + ) + + print(json.dumps(summaries, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/lazyconfig_train_net.py b/tools/lazyconfig_train_net.py index 59ae5c887a..e508d3baaa 100755 --- a/tools/lazyconfig_train_net.py +++ b/tools/lazyconfig_train_net.py @@ -79,6 +79,14 @@ def do_train(args, cfg): cfg.train.output_dir, trainer=trainer, ) + best_checkpointer = None + if comm.is_main_process() and "best_checkpointer" in cfg.train: + best_checkpointer = hooks.BestCheckpointer( + cfg.train.eval_period, + checkpointer, + **dict(cfg.train.best_checkpointer), + ) + trainer.register_hooks( [ hooks.IterationTimer(), @@ -89,6 +97,7 @@ def do_train(args, cfg): else None ), hooks.EvalHook(cfg.train.eval_period, lambda: do_test(cfg, model)), + best_checkpointer, ( hooks.PeriodicWriter( default_writers(cfg.train.output_dir, cfg.train.max_iter), diff --git a/tools/oweed_full_val_infer_annotate.py b/tools/oweed_full_val_infer_annotate.py new file mode 100644 index 0000000000..5bc94102f9 --- /dev/null +++ b/tools/oweed_full_val_infer_annotate.py @@ -0,0 +1,441 @@ +#!/usr/bin/env python +import argparse +import json +import logging +import time +from pathlib import Path + +import cv2 +import numpy as np +import torch + +from detectron2.checkpoint import DetectionCheckpointer +from detectron2.config import LazyConfig, instantiate +from detectron2.data import MetadataCatalog +from detectron2.data.detection_utils import read_image +from detectron2.data.transforms import ResizeShortestEdge +from detectron2.modeling.meta_arch import GeneralizedRCNN +from detectron2.utils.visualizer import Visualizer + + +LOGGER = logging.getLogger("oweed_full_val_infer_annotate") + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True) + parser.add_argument("--checkpoint", required=True) + parser.add_argument("--image-dir", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--score-thresh", type=float, default=0.5) + parser.add_argument("--nms-thresh", type=float, default=0.5) + parser.add_argument("--mask-thresh", type=float, default=0.5) + parser.add_argument("--warmup-iters", type=int, default=5) + parser.add_argument("--dummy-size", type=int, default=1024) + parser.add_argument( + "--postprocess-mode", + choices=("full-mask", "lowres-contour"), + default="full-mask", + help=( + "full-mask uses Detectron2's standard postprocess and materializes full-resolution " + "masks. lowres-contour keeps ROI masks compact, extracts contours at model input " + "resolution, and scales contour points to the original image." + ), + ) + parser.add_argument( + "--contours-json", + default=None, + help="Output JSON path for lowres-contour polygon predictions. Defaults inside output-dir.", + ) + parser.add_argument("--contour-alpha", type=float, default=0.45) + return parser.parse_args() + + +def gpu_mem(): + return { + "allocated_mb": torch.cuda.memory_allocated() / 1024**2, + "reserved_mb": torch.cuda.memory_reserved() / 1024**2, + "max_allocated_mb": torch.cuda.max_memory_allocated() / 1024**2, + "max_reserved_mb": torch.cuda.max_memory_reserved() / 1024**2, + } + + +def image_files(image_dir): + exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"} + return sorted(p for p in Path(image_dir).iterdir() if p.suffix.lower() in exts) + + +def run_model(model, image_rgb, resize_aug, do_postprocess=True): + original_h, original_w = image_rgb.shape[:2] + transform = resize_aug.get_transform(image_rgb) + resized = transform.apply_image(image_rgb) + tensor = torch.as_tensor(np.ascontiguousarray(resized.transpose(2, 0, 1))) + inputs = { + "image": tensor, + "height": original_h, + "width": original_w, + } + if do_postprocess: + outputs = model([inputs])[0]["instances"] + else: + outputs = model.inference([inputs], do_postprocess=False)[0] + return outputs, resized.shape[:2] + + +def timed_inference(model, image_rgb, resize_aug, do_postprocess=True): + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + start_event.record() + outputs, resized_hw = run_model(model, image_rgb, resize_aug, do_postprocess=do_postprocess) + end_event.record() + torch.cuda.synchronize() + wall_seconds = time.perf_counter() - wall_start + gpu_ms = start_event.elapsed_time(end_event) + return outputs, resized_hw, gpu_ms, wall_seconds + + +def draw_predictions(image_rgb, outputs, metadata, output_path): + vis = Visualizer(image_rgb, metadata=metadata) + drawn = vis.draw_instance_predictions(outputs.to("cpu")) + out_rgb = drawn.get_image() + out_bgr = cv2.cvtColor(out_rgb, cv2.COLOR_RGB2BGR) + cv2.imwrite(str(output_path), out_bgr) + + +def color_for_class(class_id): + # Deterministic bright-ish BGR color without depending on Visualizer internals. + palette = np.array( + [ + [230, 25, 75], + [60, 180, 75], + [255, 225, 25], + [0, 130, 200], + [245, 130, 48], + [145, 30, 180], + [70, 240, 240], + [240, 50, 230], + [210, 245, 60], + [250, 190, 190], + [0, 128, 128], + [230, 190, 255], + [170, 110, 40], + [255, 250, 200], + [128, 0, 0], + [170, 255, 195], + [128, 128, 0], + [255, 215, 180], + [0, 0, 128], + [128, 128, 128], + ], + dtype=np.uint8, + ) + rgb = palette[class_id % len(palette)] + return int(rgb[2]), int(rgb[1]), int(rgb[0]) + + +def mask_contours_from_raw_instance(mask_28, box_xyxy, resized_hw, original_hw, mask_thresh): + resized_h, resized_w = resized_hw + original_h, original_w = original_hw + + x0, y0, x1, y1 = box_xyxy.astype(float).tolist() + x0_i = max(0, min(resized_w - 1, int(np.floor(x0)))) + y0_i = max(0, min(resized_h - 1, int(np.floor(y0)))) + x1_i = max(0, min(resized_w, int(np.ceil(x1)))) + y1_i = max(0, min(resized_h, int(np.ceil(y1)))) + + crop_w = x1_i - x0_i + crop_h = y1_i - y0_i + if crop_w <= 1 or crop_h <= 1: + return [] + + resized_mask = cv2.resize(mask_28, (crop_w, crop_h), interpolation=cv2.INTER_LINEAR) + binary = (resized_mask >= mask_thresh).astype(np.uint8) + contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + scale_x = original_w / float(resized_w) + scale_y = original_h / float(resized_h) + scaled_contours = [] + for contour in contours: + if contour.shape[0] < 3: + continue + pts = contour.reshape(-1, 2).astype(np.float32) + pts[:, 0] = (pts[:, 0] + x0_i) * scale_x + pts[:, 1] = (pts[:, 1] + y0_i) * scale_y + pts[:, 0] = np.clip(pts[:, 0], 0, original_w - 1) + pts[:, 1] = np.clip(pts[:, 1], 0, original_h - 1) + scaled_contours.append(np.rint(pts).astype(np.int32)) + return scaled_contours + + +def raw_outputs_to_contours(outputs, resized_hw, original_hw, metadata, mask_thresh): + outputs_cpu = outputs.to("cpu") + boxes_resized = outputs_cpu.pred_boxes.tensor.numpy() + scores = outputs_cpu.scores.numpy() + classes = outputs_cpu.pred_classes.numpy() + masks = outputs_cpu.pred_masks[:, 0].numpy() + + resized_h, resized_w = resized_hw + original_h, original_w = original_hw + scale_x = original_w / float(resized_w) + scale_y = original_h / float(resized_h) + class_names = getattr(metadata, "thing_classes", None) + + predictions = [] + for idx, (box, score, class_id, mask) in enumerate(zip(boxes_resized, scores, classes, masks)): + contours = mask_contours_from_raw_instance( + mask, box, resized_hw, original_hw, mask_thresh=mask_thresh + ) + if not contours: + continue + + box_full = [ + float(np.clip(box[0] * scale_x, 0, original_w - 1)), + float(np.clip(box[1] * scale_y, 0, original_h - 1)), + float(np.clip(box[2] * scale_x, 0, original_w - 1)), + float(np.clip(box[3] * scale_y, 0, original_h - 1)), + ] + predictions.append( + { + "instance_index": idx, + "class_id": int(class_id), + "class_name": class_names[int(class_id)] if class_names else str(int(class_id)), + "score": float(score), + "bbox_xyxy_resized": [float(v) for v in box.tolist()], + "bbox_xyxy_fullres": box_full, + "contours_fullres": [c.reshape(-1, 2).astype(int).tolist() for c in contours], + } + ) + return predictions + + +def draw_contour_predictions(image_rgb, contour_predictions, output_path, alpha): + image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR) + overlay = image_bgr.copy() + + for pred in contour_predictions: + color = color_for_class(pred["class_id"]) + contours = [np.asarray(c, dtype=np.int32).reshape(-1, 1, 2) for c in pred["contours_fullres"]] + if not contours: + continue + cv2.fillPoly(overlay, contours, color) + cv2.polylines(image_bgr, contours, isClosed=True, color=color, thickness=2, lineType=cv2.LINE_AA) + + image_bgr = cv2.addWeighted(overlay, alpha, image_bgr, 1 - alpha, 0) + + for pred in contour_predictions: + color = color_for_class(pred["class_id"]) + x0, y0, x1, y1 = [int(round(v)) for v in pred["bbox_xyxy_fullres"]] + cv2.rectangle(image_bgr, (x0, y0), (x1, y1), color, 2, lineType=cv2.LINE_AA) + label = f'{pred["class_name"]} {pred["score"]:.2f}' + label_y = max(20, y0 - 5) + cv2.putText( + image_bgr, + label, + (x0, label_y), + cv2.FONT_HERSHEY_SIMPLEX, + 0.55, + color, + 2, + lineType=cv2.LINE_AA, + ) + + cv2.imwrite(str(output_path), image_bgr) + + +def main(): + args = parse_args() + logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)s: %(message)s") + + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + cfg = LazyConfig.load(args.config) + cfg.model._target_ = GeneralizedRCNN + cfg.model.roi_heads.box_predictor.test_score_thresh = args.score_thresh + cfg.model.roi_heads.box_predictor.test_nms_thresh = args.nms_thresh + cfg.train.init_checkpoint = args.checkpoint + + model = instantiate(cfg.model) + model.to(cfg.train.device) + model.eval() + DetectionCheckpointer(model).load(args.checkpoint) + + dataset_name = cfg.dataloader.evaluator.dataset_name + metadata = MetadataCatalog.get(dataset_name) + resize_aug = ResizeShortestEdge(short_edge_length=1024, max_size=1024) + + files = image_files(args.image_dir) + if not files: + raise RuntimeError(f"No images found in {args.image_dir}") + + run_log = { + "config": args.config, + "checkpoint": args.checkpoint, + "image_dir": args.image_dir, + "output_dir": str(out_dir), + "score_thresh": args.score_thresh, + "nms_thresh": args.nms_thresh, + "mask_thresh": args.mask_thresh, + "postprocess_mode": args.postprocess_mode, + "resize_short_edge": 1024, + "resize_max_size": 1024, + "num_images": len(files), + "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, + "warmup": {}, + "images": [], + } + contour_records = { + "format": "oweed_scaled_lowres_contours_v1", + "description": ( + "Contours are extracted from compact ROI masks at resized model-input scale, " + "then scaled to original image coordinates. Dense full-resolution masks are not stored." + ), + "config": args.config, + "checkpoint": args.checkpoint, + "image_dir": args.image_dir, + "score_thresh": args.score_thresh, + "nms_thresh": args.nms_thresh, + "mask_thresh": args.mask_thresh, + "resize_short_edge": 1024, + "resize_max_size": 1024, + "images": [], + } + contours_json_path = ( + Path(args.contours_json) + if args.contours_json + else out_dir / "contours_fullres_scaled_from_lowres.json" + ) + do_postprocess = args.postprocess_mode == "full-mask" + + torch.cuda.reset_peak_memory_stats() + with torch.no_grad(): + dummy = torch.zeros( + 3, args.dummy_size, args.dummy_size, dtype=torch.uint8, device="cpu" + ).numpy() + dummy_rgb = np.transpose(dummy, (1, 2, 0)) + warmup_times = [] + for _ in range(args.warmup_iters): + _, _, gpu_ms, wall_seconds = timed_inference( + model, dummy_rgb, resize_aug, do_postprocess=do_postprocess + ) + warmup_times.append({"gpu_ms": gpu_ms, "wall_seconds": wall_seconds}) + run_log["warmup"] = { + "iters": args.warmup_iters, + "dummy_hwc": [args.dummy_size, args.dummy_size, 3], + "times": warmup_times, + "memory_after_warmup": gpu_mem(), + } + + torch.cuda.reset_peak_memory_stats() + total_gpu_ms = 0.0 + total_wall_seconds = 0.0 + total_contour_seconds = 0.0 + total_draw_seconds = 0.0 + max_detections = 0 + for index, path in enumerate(files, start=1): + image_rgb = read_image(str(path), format="RGB") + torch.cuda.reset_peak_memory_stats() + outputs, resized_hw, gpu_ms, wall_seconds = timed_inference( + model, image_rgb, resize_aug, do_postprocess=do_postprocess + ) + num_instances = len(outputs) + max_detections = max(max_detections, num_instances) + total_gpu_ms += gpu_ms + total_wall_seconds += wall_seconds + + output_path = ( + out_dir / f"{path.stem}_pred_conf{args.score_thresh:.2f}_nms{args.nms_thresh:.2f}.jpg" + ) + contour_seconds = 0.0 + draw_seconds = 0.0 + kept_contours = None + if args.postprocess_mode == "full-mask": + draw_start = time.perf_counter() + draw_predictions(image_rgb, outputs, metadata, output_path) + draw_seconds = time.perf_counter() - draw_start + else: + contour_start = time.perf_counter() + contour_predictions = raw_outputs_to_contours( + outputs, + resized_hw, + image_rgb.shape[:2], + metadata, + mask_thresh=args.mask_thresh, + ) + contour_seconds = time.perf_counter() - contour_start + kept_contours = len(contour_predictions) + + draw_start = time.perf_counter() + draw_contour_predictions( + image_rgb, + contour_predictions, + output_path, + alpha=args.contour_alpha, + ) + draw_seconds = time.perf_counter() - draw_start + + contour_records["images"].append( + { + "file": str(path), + "output": str(output_path), + "original_hw": list(image_rgb.shape[:2]), + "resized_hw": list(resized_hw), + "num_raw_detections": num_instances, + "num_contour_predictions": kept_contours, + "predictions": contour_predictions, + } + ) + + total_contour_seconds += contour_seconds + total_draw_seconds += draw_seconds + + record = { + "index": index, + "file": str(path), + "output": str(output_path), + "original_hw": list(image_rgb.shape[:2]), + "resized_hw": list(resized_hw), + "detections": num_instances, + "contour_predictions": kept_contours, + "gpu_ms": gpu_ms, + "wall_seconds": wall_seconds, + "contour_seconds": contour_seconds, + "draw_seconds": draw_seconds, + "memory": gpu_mem(), + } + run_log["images"].append(record) + LOGGER.info( + "%03d/%03d %s det=%d contours=%s gpu_ms=%.2f wall=%.3fs peak_alloc=%.1fMB", + index, + len(files), + path.name, + num_instances, + "-" if kept_contours is None else kept_contours, + gpu_ms, + wall_seconds, + record["memory"]["max_allocated_mb"], + ) + + run_log["summary"] = { + "total_gpu_seconds": total_gpu_ms / 1000.0, + "total_wall_seconds_model_only": total_wall_seconds, + "avg_gpu_ms_per_image": total_gpu_ms / len(files), + "avg_wall_seconds_model_only": total_wall_seconds / len(files), + "total_contour_seconds": total_contour_seconds, + "avg_contour_seconds": total_contour_seconds / len(files), + "total_draw_seconds": total_draw_seconds, + "avg_draw_seconds": total_draw_seconds / len(files), + "max_detections_per_image": max_detections, + "final_memory": gpu_mem(), + } + log_path = out_dir / "inference_timing_gpu_mem.json" + log_path.write_text(json.dumps(run_log, indent=2)) + if args.postprocess_mode == "lowres-contour": + contours_json_path.write_text(json.dumps(contour_records, indent=2)) + LOGGER.info("Wrote scaled contour predictions to %s", contours_json_path) + LOGGER.info("Wrote timing/memory log to %s", log_path) + + +if __name__ == "__main__": + main()