Bound proposal-filter inference memory

This commit is contained in:
Jens
2026-07-27 08:15:06 +02:00
parent d8fc66230e
commit 81283551ad
+23 -12
View File
@@ -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,