ultralytics 8.0.203 DDP argparse and Tracker fixes (#6007)

This commit is contained in:
Glenn Jocher 2023-10-29 21:14:50 +01:00 committed by GitHub
parent ab0b47e386
commit 465df3024f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 50 additions and 22 deletions

View file

@ -1,6 +1,7 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
from functools import partial
from pathlib import Path
import torch
@ -40,8 +41,12 @@ def on_predict_start(predictor, persist=False):
def on_predict_postprocess_end(predictor):
"""Postprocess detected boxes and update with object tracking."""
bs = predictor.dataset.bs
im0s = predictor.batch[1]
path, im0s = predictor.batch[:2]
for i in range(bs):
if predictor.vid_path[i] != str(predictor.save_dir / Path(path[i]).name): # new video
predictor.trackers[i].reset()
det = predictor.results[i].boxes.cpu().numpy()
if len(det) == 0:
continue