Files
ieemoo-ai-imageassessment/ytracking/track_.py.bak
2024-11-27 15:37:10 +08:00

227 lines
9.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# 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()