update
This commit is contained in:
lee
226
ytracking/track_.py.bak
Normal file
226
ytracking/track_.py.bak
Normal file
@ -0,0 +1,226 @@
|
||||
# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
|
||||
import argparse
|
||||
import csv
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import glob
|
||||
import numpy as np
|
||||
import pickle
|
||||
import torch
|
||||
|
||||
# =============================================================================
|
||||
# FILE = Path(__file__).resolve()
|
||||
# ROOT = FILE.parents[0] # YOLOv5 root directory
|
||||
# if str(ROOT) not in sys.path:
|
||||
# sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||
# ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
|
||||
# =============================================================================
|
||||
sys.path.append('./ytracking')
|
||||
from models.common import DetectMultiBackend
|
||||
from utils.dataloaders import LoadImages
|
||||
from utils.general import (LOGGER, Profile, check_img_size, check_requirements, colorstr, cv2,
|
||||
increment_path, non_max_suppression, scale_boxes, strip_optimizer)
|
||||
from utils.torch_utils import select_device, smart_inference_mode
|
||||
|
||||
'''集成跟踪模块，输出跟踪结果文件 .npy'''
|
||||
# from ultralytics.engine.results import Boxes # Results
|
||||
# from ultralytics.utils import IterableSimpleNamespace, yaml_load
|
||||
from tracking.utils.plotting import Annotator, colors
|
||||
from tracking.utils import Boxes, IterableSimpleNamespace, yaml_load
|
||||
from tracking.trackers import BOTSORT, BYTETracker
|
||||
from tracking.utils.showtrack import drawtracks
|
||||
|
||||
|
||||
# tracker_yaml = r"./tracking/trackers/cfg/botsort.yaml"
|
||||
|
||||
def init_trackers(tracker_yaml=None, bs=1):
|
||||
"""
|
||||
Initialize trackers for object tracking during prediction.
|
||||
"""
|
||||
# tracker_yaml = r"./tracking/trackers/cfg/botsort.yaml"
|
||||
|
||||
TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}
|
||||
|
||||
cfg = IterableSimpleNamespace(**yaml_load(tracker_yaml))
|
||||
trackers = []
|
||||
for _ in range(bs):
|
||||
tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)
|
||||
trackers.append(tracker)
|
||||
|
||||
return trackers
|
||||
|
||||
|
||||
@smart_inference_mode()
|
||||
def run(
|
||||
weights=r"D:/Project/ieemoo-ai/tools/ckpts/best_158734_cls11_noaug10.pt", # model path or triton URL
|
||||
source=r"D:/Project/ieemoo-ai/testdata/88.mp4", # file/dir/URL/glob/screen/0(webcam)
|
||||
|
||||
project=r'./runs/detect', # save results to project/name
|
||||
name='exp', # save results to project/name
|
||||
|
||||
tracker_yaml="D:/Project/ieemoo-ai/ytracking/tracking/trackers/cfg/botsort.yaml",
|
||||
imgsz=(640, 640), # inference size (height, width)
|
||||
conf_thres=0.25, # confidence threshold
|
||||
iou_thres=0.45, # NMS IOU threshold
|
||||
max_det=1000, # maximum detections per image
|
||||
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
||||
bs=1, # batch_size
|
||||
save_txt=False, # save results to *.txt
|
||||
save_img=True, # do not save images/videos
|
||||
classes=None, # filter by class: --class 0, or --class 0 2 3
|
||||
agnostic_nms=False, # class-agnostic NMS
|
||||
augment=False, # augmented inference
|
||||
visualize=False, # visualize features
|
||||
update=False, # update all models
|
||||
exist_ok=False, # existing project/name ok, do not increment
|
||||
line_thickness=3, # bounding box thickness (pixels)
|
||||
hide_labels=False, # hide labels
|
||||
hide_conf=False, # hide confidencesL
|
||||
half=False, # use FP16 half-precision inference
|
||||
dnn=False, # use OpenCV DNN for ONNX inference
|
||||
vid_stride=1, # video frame-rate stride
|
||||
):
|
||||
# Load model
|
||||
device = select_device(device)
|
||||
model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
|
||||
stride, names, pt = model.stride, model.names, model.pt
|
||||
imgsz = check_img_size(imgsz, s=stride) # check image size
|
||||
|
||||
# Run inference
|
||||
model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz)) # warmup
|
||||
|
||||
##=============================生成文件夹 save_dir，存储检测跟踪图像
|
||||
source = str(source)
|
||||
save_dir = Path(project) / Path(source).stem
|
||||
if save_dir.exists():
|
||||
print(Path(source).stem)
|
||||
# return
|
||||
|
||||
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
|
||||
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
|
||||
else:
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Dataloader
|
||||
seen, dt = 0, (Profile(), Profile(), Profile())
|
||||
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
|
||||
|
||||
## ================================================= 生成跟踪器对象
|
||||
tracker = init_trackers(tracker_yaml, bs)[0]
|
||||
track_boxes = np.empty((0, 9), dtype=np.float32)
|
||||
features_dict = {}
|
||||
|
||||
for path, im, im0s, vid_cap, s in dataset:
|
||||
# img preprocess
|
||||
with dt[0]:
|
||||
im = torch.from_numpy(im).to(model.device)
|
||||
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
|
||||
im /= 255 # 0 - 255 to 0.0 - 1.0
|
||||
if len(im.shape) == 3:
|
||||
im = im[None] # expand for batch dim
|
||||
# Inference
|
||||
with dt[1]:
|
||||
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
|
||||
pred = model(im, augment=augment, visualize=visualize)
|
||||
# NMS
|
||||
with dt[2]:
|
||||
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
|
||||
|
||||
# Process predictions
|
||||
for i, det in enumerate(pred): # per image
|
||||
seen += 1
|
||||
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
|
||||
|
||||
im0_ant = im0.copy()
|
||||
|
||||
p = Path(p) # to Path
|
||||
save_path = str(save_dir / p.name) # im.jpg
|
||||
s += '%gx%g ' % im.shape[2:] # print string
|
||||
|
||||
annotator = Annotator(im0_ant, line_width=line_thickness, example=str(names)) if save_img else None
|
||||
if len(det):
|
||||
# Rescale boxes from img_size to im0 size
|
||||
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
|
||||
|
||||
# boxes_and_imgs.append((det.cpu().numpy(), im0, frame))
|
||||
## ================================================================ writed by WQG
|
||||
det_tracking = Boxes(det, im0.shape).cpu().numpy()
|
||||
tracks = tracker.update(det_tracking, im0)
|
||||
|
||||
if len(tracks):
|
||||
track_boxes = np.concatenate([track_boxes, tracks], axis=0)
|
||||
feat_dict = {int(x.idx): x.curr_feat for x in tracker.tracked_stracks if x.is_activated}
|
||||
frame_id = track_boxes[0, 7]
|
||||
features_dict.update({int(frame_id): feat_dict})
|
||||
|
||||
if annotator is not None:
|
||||
for *xyxy, id, conf, cls, fid, bid in reversed(tracks):
|
||||
name = ('' if id == -1 else f'id:{int(id)} ') + names[int(cls)]
|
||||
label = None if hide_labels else (name if hide_conf else f'{name} {conf:.2f}')
|
||||
|
||||
if id >= 0 and cls == 0:
|
||||
color = colors(int(cls), True)
|
||||
elif id >= 0 and cls != 0:
|
||||
color = colors(int(id), True)
|
||||
else:
|
||||
color = colors(19, True) # 19为调色板的最后一个元素
|
||||
annotator.box_label(xyxy, label, color=color)
|
||||
|
||||
# Save tracking image
|
||||
if annotator is not None:
|
||||
save_path_img, ext = os.path.splitext(save_path)
|
||||
imgpath = save_path_img + f"_{dataset.frame}.png"
|
||||
cv2.imwrite(Path(imgpath), annotator.result())
|
||||
|
||||
# Print time (inference-only)
|
||||
LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
|
||||
|
||||
## ======================================================================== written by WQG
|
||||
''' track_boxes: Array, [x1, y1, x2, y2, track_id, score, cls, frame_index, box_id] '''
|
||||
if save_img:
|
||||
filename = os.path.split(save_path_img)[-1]
|
||||
'''====== save in './run/detect/' ======'''
|
||||
imgshow = drawtracks(track_boxes)
|
||||
showpath_1 = save_path_img + "_show.png"
|
||||
cv2.imwrite(Path(showpath_1), imgshow)
|
||||
|
||||
'''====== save tracks data ======'''
|
||||
tracks_dir = Path('D:/Project/ieemoo-ai/ytracking/tracking/tracking/data/tracks/')
|
||||
if not tracks_dir.exists():
|
||||
tracks_dir.mkdir(parents=True, exist_ok=True)
|
||||
tracks_path = tracks_dir.joinpath(filename + ".npy")
|
||||
np.save(tracks_path, track_boxes)
|
||||
|
||||
'''====== save reid features data ======'''
|
||||
feats_dir = Path('D:/Project/ieemoo-ai/ytracking/tracking/data/trackfeats/')
|
||||
if not feats_dir.exists():
|
||||
feats_dir.mkdir(parents=True, exist_ok=True)
|
||||
feats_path = feats_dir.joinpath(f'{filename}.pkl')
|
||||
with open(feats_path, 'wb') as file:
|
||||
pickle.dump(features_dict, file)
|
||||
|
||||
# Print results
|
||||
t = tuple(x.t / seen * 1E3 for x in dt) # speeds per image
|
||||
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
|
||||
if save_txt or save_img:
|
||||
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
|
||||
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
|
||||
if update:
|
||||
strip_optimizer(weights[0]) # update model (to fix SourceChangeWarning)
|
||||
|
||||
|
||||
def main():
|
||||
ROOT = Path(Path.cwd())
|
||||
check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
|
||||
|
||||
optdict = {'weights': r"D:/Project/ieemoo-ai/tools/ckpts/best_158734_cls11_noaug10.pt",
|
||||
'source': r"D:/Project/ieemoo-ai/testdata/88.mp4",
|
||||
}
|
||||
|
||||
run(**optdict)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user