From 81283551ad6185e79c482b72e472409de7043648 Mon Sep 17 00:00:00 2001 From: Jens Date: Mon, 27 Jul 2026 08:15:06 +0200 Subject: [PATCH] Bound proposal-filter inference memory --- scripts/train_building_proposal_filter.py | 35 +++++++++++++++-------- 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/scripts/train_building_proposal_filter.py b/scripts/train_building_proposal_filter.py index eaf2cb43..de90aac7 100644 --- a/scripts/train_building_proposal_filter.py +++ b/scripts/train_building_proposal_filter.py @@ -84,8 +84,8 @@ def materialize_split( for label, boxes in (("positive", positives), ("negative", negatives)): directory = output_dir / split / label directory.mkdir(parents=True, exist_ok=True) - for box in boxes: - path = directory / f"{tile['sample_slug']}_{tile['tile_index']:04d}_{counts[label]:07d}.png" + for box_index, box in enumerate(boxes): + path = directory / f"{Path(tile['image_path']).stem}_{label}_{box_index:04d}.png" expanded_crop(image, box, crop_scale).save(path) counts[label] += 1 return counts @@ -121,6 +121,7 @@ def main() -> int: parser.add_argument("--negative-match-iou", type=float, default=0.05) parser.add_argument("--epochs", type=int, default=8) parser.add_argument("--batch", type=int, default=64) + parser.add_argument("--proposal-chunk-size", type=int, default=16) parser.add_argument("--seed", type=int, default=20260727) parser.add_argument("--force", action="store_true") args = parser.parse_args() @@ -143,18 +144,27 @@ def main() -> int: for split in ("train", "val") } proposal_model = YOLO(str(args.proposal_model)) + if args.proposal_chunk_size < 1: + raise SystemExit("--proposal-chunk-size must be positive") counts = {} for split, tiles in split_tiles.items(): - results = proposal_model.predict( - [tile["image_path"] for tile in tiles], conf=args.proposal_confidence, imgsz=640, - max_det=1000, device=args.device, verbose=False, - ) - counts[split] = materialize_split( - tiles=tiles, results=results, output_dir=args.output_dir / "crops", split=split, - crop_scale=args.crop_scale, max_positives_per_tile=args.max_positives_per_tile, - max_negatives_per_tile=args.max_negatives_per_tile, negative_match_iou=args.negative_match_iou, - ) - del results + counts[split] = {"positive": 0, "negative": 0} + for start in range(0, len(tiles), args.proposal_chunk_size): + chunk = tiles[start : start + args.proposal_chunk_size] + results = proposal_model.predict( + [tile["image_path"] for tile in chunk], conf=args.proposal_confidence, imgsz=640, + max_det=1000, device=args.device, verbose=False, + ) + chunk_counts = materialize_split( + tiles=chunk, results=results, output_dir=args.output_dir / "crops", split=split, + crop_scale=args.crop_scale, max_positives_per_tile=args.max_positives_per_tile, + max_negatives_per_tile=args.max_negatives_per_tile, negative_match_iou=args.negative_match_iou, + ) + for label, count in chunk_counts.items(): + counts[split][label] += count + del results + if torch.cuda.is_available(): + torch.cuda.empty_cache() del proposal_model if torch.cuda.is_available(): @@ -206,6 +216,7 @@ def main() -> int: "summary": str(args.summary), "summary_sha256": sha256(args.summary), "proposal_model": str(args.proposal_model), "proposal_model_sha256": sha256(args.proposal_model), "device": args.device, "proposal_confidence": args.proposal_confidence, "crop_scale": args.crop_scale, + "proposal_chunk_size": args.proposal_chunk_size, "negative_match_iou": args.negative_match_iou, "counts": counts, "epochs": args.epochs, "history": history, "selected_threshold_source": "validation_only", "selected_threshold": threshold, "validation_metrics": validation_metrics,