INNER CODE UNIT · Python

add_balls_before_SORT

chonyy/ML-auto-baseball-pitching-overlay · src/get_pitch_frames.py:198

def add_balls_before_SORT(frames, detected, tracked, tracker_min_hits):
    distance_threshold = 100
    first_ball = tracked[0]
    color = first_ball[2]
    balls_to_add = []

    # Get the untracked balls that's close enough to the first tracked ball
    for untracked in detected[-(tracker_min_hits + 1) :]:
        if distance(untracked, first_ball) < distance_threshold:
            untracked.append(color)
            balls_to_add.append(untracked)

    # Add the untracked balls to frame
    modify_frames = frames[-(tracker_min_hits + 1) :]
    balls_to_add_temp = copy.deepcopy(balls_to_add)

    for point in balls_to_add_temp:
        del point[2]

View source record →