Я предоставляю вам измененный код. Буду рад, если вы мне немного поможете. Я также изменил функцию обновления Strongsort.py
track_v7.py
Код: Выделить всё
import os
import sys
import argparse
import time
from pathlib import Path
import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
import numpy as np
import zmq
import json
import time
import traceback
import warnings
warnings.filterwarnings('ignore')
# Constants
ROOT = Path(__file__).resolve().parents[0]
WEIGHTS = ROOT / 'weights'
from strong_sort.utils.parser import get_config
from strong_sort.strong_sort import StrongSORT
from yolov7.utils.datasets import LoadStreams, LoadImages
from yolov7.utils.general import check_imshow
from yolov7.utils.torch_utils import select_device
from yolov7.utils.plots import plot_one_box
def process_frame_and_boxes(strong, names, colors, device, frame_np, bounding_boxes):
try:
# Prepare frame
frame = cv2.imdecode(frame_np, cv2.IMREAD_COLOR)
im0 = frame.copy()
# Initialize variables
nr_sources = len(bounding_boxes)
outputs = [None] * nr_sources
trajectory = {}
# Process each bounding box
for i, bbox in enumerate(bounding_boxes):
xmin = bbox["xmin"]
ymin = bbox["ymin"]
xmax = bbox["xmax"]
ymax = bbox["ymax"]
cls_str = bbox["class"]
conf = bbox["probability"]
width = xmax - xmin
height = ymax - ymin
#xmin, ymin, xmax, ymax, cls, conf = bbox
try:
cls = names.index(cls_str)
except ValueError:
print(f"Unknown class: {cls_str}")
continue
# Prepare bounding box data
bboxes = np.array([xmin, ymin, width, height])
#id = i # Assuming id can be derived from the loop index
cls = int(cls)
conf = float(conf)
# Track object using StrongSORT
xywhs = torch.Tensor([bboxes]) # Convert to tensor
confs = torch.Tensor([conf])
clss = torch.Tensor([cls])
print("BB: " ,bbox)
try:
outputs[i] = strong.update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
print(f"Box:{xywhs} confidence:{confs} output:{outputs[i]}")
except Exception as e:
print(f"Unable to process properly: {e}")
#outputs[i] = strongsort_list[i].update(xywhs, confs, clss, im0)
# draw boxes for visualization
if len(outputs[i]) > 0:
for j, (output, conf) in enumerate(zip(outputs[i], confs)):
bboxes = output[0:4]
id = output[4]
cls = int(output[5])
# Draw trajectory if required
if draw:
center = ((int(bboxes[0]) + int(bboxes[2])) // 2, (int(bboxes[1]) + int(bboxes[3])) // 2)
if id not in trajectory:
trajectory[id] = []
trajectory[id].append(center)
for i1 in range(1, len(trajectory[id])):
if trajectory[id][i1 - 1] is None or trajectory[id][i1] is None:
continue
thickness = 2
try:
cv2.line(im0, trajectory[id][i1 - 1], trajectory[id][i1], (0, 0, 255), thickness)
except Exception as e:
print(f"Error drawing trajectory: {e}")
# Draw bounding boxes on image
if show_vid:
label = f"{names[cls]} ID: {id}"
plot_one_box(bboxes, im0, label=label, color=colors[cls], line_thickness=line_thickness)
# Display or save images/videos if required
if show_vid:
cv2.imshow("Tracking", im0)
cv2.waitKey(1) # Adjust as necessary for frame rate
#if save_img:
# cv2.imwrite(save_path, im0)
#if save_vid:
# vid_writer.write(im0)'''
except Exception as e:
print(f"Error: {e}")
print(traceback.format_exc())
def set_logging():
import logging
logging.basicConfig(format="%(message)s", level=logging.INFO)
def track_bounding_boxes():
# Process frames and boxes
context = zmq.Context()
socket = context.socket(zmq.SUB)
socket.connect("tcp://127.0.0.1:5555")
socket.setsockopt_string(zmq.SUBSCRIBE, "")
# Initialize StrongSORT
cfg = get_config()
cfg.merge_from_file(opt.config_strongsort)
# Placeholder for class names, should be provided
names = ['person', 'vehicle', 'wheel', 'weapon']
# Placeholder for colors, should be provided
colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
#nr_sources = 1
# Create as many StrongSORT instances as there are sources
#strongsort_list = []
#for i in range(nr_sources):
# strongsort_list.append(
strong = StrongSORT(
strong_sort_weights,
device,
max_dist=cfg.STRONGSORT.MAX_DIST,
max_iou_distance=cfg.STRONGSORT.MAX_IOU_DISTANCE,
max_age=cfg.STRONGSORT.MAX_AGE,
n_init=cfg.STRONGSORT.N_INIT,
nn_budget=cfg.STRONGSORT.NN_BUDGET,
mc_lambda=cfg.STRONGSORT.MC_LAMBDA,
ema_alpha=cfg.STRONGSORT.EMA_ALPHA,
)
#)
try:
while True:
print("Listening...")
msg = socket.recv_string()
frame_data = socket.recv()
frame_np = np.frombuffer(frame_data, np.uint8)
if msg and frame_np.size > 0:
bounding_boxes = json.loads(msg)
print("bounding boxes: ", bounding_boxes)
process_frame_and_boxes(strong, names, colors, device, frame_np, bounding_boxes)
else:
print("Error: Empty message received")
except KeyboardInterrupt:
pass
except Exception as e:
print(f"Error during processing: {e}")
finally:
socket.close()
context.term()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--strong-sort-weights', type=str, default=WEIGHTS / 'osnet_x0_25_msmt17.pt')
parser.add_argument('--config-strongsort', type=str, default='strong_sort/configs/strong_sort.yaml')
parser.add_argument('--save-img', action='store_true', help='save results to *.jpg')
parser.add_argument('--save-vid', action='store_true', help='save video tracking results')
parser.add_argument('--show-vid', action='store_true', default=True , help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--project', default='runs/track', help='save results to project/name')
parser.add_argument('--exp-name', default='exp', help='save results to project/name')
parser.add_argument('--line-thickness', default=1, type=int, help='bounding box thickness (pixels)')
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
parser.add_argument('--hide-class', default=False, action='store_true', help='hide IDs')
parser.add_argument('--draw', action='store_true', help='display object trajectory lines')
opt = parser.parse_args()
set_logging()
strong_sort_weights = opt.strong_sort_weights
config_strongsort = opt.config_strongsort
save_img = opt.save_img
save_vid = opt.save_vid
show_vid = opt.show_vid
save_txt = opt.save_txt
line_thickness = opt.line_thickness
draw = opt.draw
save_dir = Path(opt.project) / opt.exp_name
save_dir.mkdir(parents=True, exist_ok=True)
# Initialize variables
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Update as per your implementation
track_bounding_boxes()
Код: Выделить всё
def update(self, bbox_xywh, confidences, classes, ori_img):
try:
print(f"Box_xywhs: {bbox_xywh} confdence:{confidences}")
self.height, self.width = ori_img.shape[:2]
# generate detections
features = self._get_features(bbox_xywh, ori_img)
bbox_tlwh = bbox_xywh # self._xywh_to_tlwh(bbox_xywh)
detections = [Detection(bbox_tlwh[i], conf, features[i]) for i, conf in enumerate(
confidences)]
# run on non-maximum supression
boxes = np.array([d.tlwh for d in detections])
scores = np.array([d.confidence for d in detections])
for det in detections:
print(f"tlwh= {det.tlwh} confidence= {det.confidence}")
print("Classes: ", classes)
print("Confidence: ", confidences)
print("Update_boxes: ", boxes)
# update tracker
self.tracker.predict()
self.tracker.update(detections, classes, confidences)
# output bbox identities
outputs = []
for track in self.tracker.tracks:
if not track.is_confirmed() or track.time_since_update > 1:
continue
box = track.to_tlwh()
print("U11111111")
x1, y1, x2, y2 = self._tlwh_to_xyxy(box)
track_id = track.track_id
class_id = track.class_id
conf = track.conf
outputs.append(np.array([x1, y1, x2, y2, track_id, class_id, conf]))
#print("Converted values:", outputs)
if len(outputs) > 0:
outputs = np.stack(outputs, axis=0)
return outputs
except Exception as e:
print(f"Error in strongsort: {e}")
print(traceback.format_exc())
- Невозможно получить данные отслеживания
- если да, то невозможно отслеживать несколько ограничивающих рамок.
- Визуализация очень сильно тормозит.
Спасибо
Подробнее здесь: https://stackoverflow.com/questions/786 ... e-tracking
