lichen/detecttracking
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

modify at output data format

Browse Source
This commit is contained in:
王庆刚
2024-12-10 19:01:54 +08:00
parent 1e6c5deee4
commit afd033b965
15 changed files with 545 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
track_reid.py
View File

@ -196,8 +196,8 @@ def yolo_resnet_tracker(
tracker = init_trackers(tracker_yaml, bs)[0]
dt = (Profile(), Profile(), Profile())
track_boxes = np.empty((0, 9), dtype = np.float32)
TracksDict = {}
# trackerBoxes = np.empty((0, 9), dtype = np.float32)
yoloResnetTracker = []
for path, im, im0s, vid_cap, s in dataset:
with dt[0]:
im = torch.from_numpy(im).to(model.device)
@ -245,24 +245,23 @@ def yolo_resnet_tracker(
tracks[:, 7] = frameId
# trackerBoxes = np.concatenate([trackerBoxes, tracks], axis=0)
'''================== 1. 存储 dets/subimgs/features Dict ============='''
imgs, features = ReIDEncoder.inference(im0, tracks)
# imgs, features = inference_image(im0, tracks)
# TrackerFeats = np.concatenate([TrackerFeats, features], axis=0)
imgdict = {}
boxdict = {}
featdict = {}
for ii, bid in enumerate(tracks[:, 8]):
imgdict.update({int(bid): imgs[ii]}) # [f"img_{int(bid)}"] = imgs[i]
boxdict.update({int(bid): tracks[ii, :]}) # [f"box_{int(bid)}"] = tracks[i, :]
featdict.update({int(bid): features[ii, :]}) # [f"feat_{int(bid)}"] = features[i, :]
TracksDict[f"frame_{int(frameId)}"] = {"imgs":imgdict, "boxes":boxdict, "feats":featdict}
track_boxes = np.concatenate([track_boxes, tracks], axis=0)
featdict.update({f"{int(frameId)}_{int(bid)}": features[ii, :]}) # [f"feat_{int(bid)}"] = features[i, :]
frameDict = {"path": path,
"fid": int(frameId),
"bboxes": det,
"tboxes": tracks,
"imgs": imgs,
"feats": featdict}
yoloResnetTracker.append(frameDict)
# imgs, features = inference_image(im0, tracks)
# TrackerFeats = np.concatenate([TrackerFeats, features], axis=0)
'''================== 2. 提取手势位置 ==================='''
for *xyxy, id, conf, cls, fid, bid in reversed(tracks):
@ -314,11 +313,8 @@ def yolo_resnet_tracker(
# Print time (inference-only)
LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
## track_boxes: Array, [x1, y1, x2, y2, track_id, score, cls, frame_index, box_id]
TracksDict.update({"TrackBoxes": track_boxes})
return TracksDict
return yoloResnetTracker