update

ytracking/track_.py

@@ -0,0 +1,159 @@
+# 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
+from tools.config import cfg
+
+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
+import time
+
+
+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=cfg.tracking_model,  # model path or triton URL
+        Model,  # model path or triton URL
+        source=None,  # file/dir/URL/glob/screen/0(webcam)
+
+        project=r'./runs/detect',  # save results to project/name
+
+        tracker_yaml=cfg.botsort,
+        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_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
+        line_thickness=3,  # bounding box thickness (pixels)
+        half=False,  # use FP16 half-precision inference
+        dnn=False,  # use OpenCV DNN for ONNX inference
+        vid_stride=1,  # video frame-rate stride
+):
+    if source is None:
+        raise ValueError("Have to provide --source argument")
+
+    # Load model
+    # device = select_device(device)
+    # model = DetectMultiBackend(weights, device=device, dnn=dnn, fp16=half)
+    if Model is None:
+        raise ValueError("Have to provide --model argument")
+    model = Model.yoloModel
+    print(model.stride, model.names, model.pt)
+    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
+
+    # 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 = {}
+    frameid_img = {}
+
+    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
+            frameid_img[seen] = im0s.copy()
+
+            p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
+
+            s += '%gx%g ' % im.shape[2:]  # print string
+
+            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):
+                    tracks[:, 7] = seen
+                    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 = tracks[0, 7]
+                    features_dict.update({int(frame_id): feat_dict})
+
+    return track_boxes, features_dict, frameid_img
+
+
+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()