Not able to acquire Tracking.

[federista]

I have modified your script to work on custom weights. I am using yolov4 for detection. I am only using the tracking algorithm for your repository and providing the image source and all the bounding box information from my yolov4 node using socket.

Problem

1. Unable to acquire tracking output
2. if yes, then unable to tracking multiple bounding boxes
3. The visualization lags really bad.

I am providing you the modified code. Will be glad if you can help me a bit. I have also modified the update function of 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()

Strongsort.py (Update function)

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

Would highly appreciate if you can guide me a bit with it. My concepts are not that strong so sorry if you notice some silly mistakes
Thank you