cjm-byte-track

A standalone Python implementation of the ByteTrack multi-object tracker based on the official implementation.

A standalone Python implementation of the ByteTrack multi-object tracker based on the official implementation.

Install

pip install cjm_byte_track

Tutorial:

How to use

# Import ByteTrack package
from cjm_byte_track.core import BYTETracker
from cjm_byte_track.matching import match_detections_with_tracks
# Initialize a ByteTracker object
tracker = BYTETracker(track_thresh=0.25, track_buffer=30, match_thresh=0.8, frame_rate=frame_fps)

with tqdm(total=frames, desc="Processing frames") as pbar:
    while video_capture.isOpened():
        ret, frame = video_capture.read()
        if ret:
            
            # Prepare an input image for inference
            rgb_img, input_dims, offsets, min_img_scale, input_img = prepare_image_for_inference(frame, test_sz, max_stride)
                        
            # Convert the existing input image to NumPy format
            input_tensor_np = np.array(input_img, dtype=np.float32).transpose((2, 0, 1))[None]/255

            # Start performance counter`m
            start_time = time.perf_counter()
                        
            # Run inference
            outputs = session.run(None, {"input": input_tensor_np})[0]

            # Process the model output
            proposals = process_outputs(outputs, input_tensor_np.shape[input_dim_slice], bbox_conf_thresh)
            
            # Apply non-max suppression to the proposals with the specified threshold
            proposal_indices = nms_sorted_boxes(calc_iou(proposals[:, :-2]), iou_thresh)
            proposals = proposals[proposal_indices]
            
            bbox_list = (proposals[:,:4]+[*offsets, 0, 0])*min_img_scale
            label_list = [class_names[int(idx)] for idx in proposals[:,4]]
            probs_list = proposals[:,5]

            # Update tracker with detections.
            track_ids = [-1]*len(bbox_list)

            # Convert to tlbr format
            tlbr_boxes = bbox_list.copy()
            tlbr_boxes[:, 2:4] += tlbr_boxes[:, :2]

            # Update tracker with detections
            tracks = tracker.update(
                output_results=np.concatenate([tlbr_boxes, probs_list[:, np.newaxis]], axis=1),
                img_info=rgb_img.size,
                img_size=rgb_img.size)
            track_ids = match_detections_with_tracks(tlbr_boxes=tlbr_boxes, track_ids=track_ids, tracks=tracks)

            # End performance counter
            end_time = time.perf_counter()
            # Calculate the combined FPS for object detection and tracking
            fps = 1 / (end_time - start_time)
            # Display the frame rate in the progress bar
            pbar.set_postfix(fps=fps)

            # Filter object detections based on tracking results
            bbox_list, label_list, probs_list, track_ids = zip(*[(bbox, label, prob, track_id) 
                                                                 for bbox, label, prob, track_id 
                                                                 in zip(bbox_list, label_list, probs_list, track_ids) if track_id != -1])

            # Annotate the current frame with bounding boxes and tracking IDs
            annotated_img = draw_bboxes_pil(
                image=rgb_img, 
                boxes=bbox_list, 
                labels=[f"{track_id}-{label}" for track_id, label in zip(track_ids, label_list)],
                probs=probs_list,
                colors=[int_colors[class_names.index(i)] for i in label_list],  
                font=font_file,
            )
            annotated_frame = cv2.cvtColor(np.array(annotated_img), cv2.COLOR_RGB2BGR)
            
            video_writer.write(annotated_frame)
            pbar.update(1)
        else:
            break
video_capture.release()
video_writer.release()