A-Raj/arrow-tracking

1import cv22import numpy as np3from ultralytics import YOLO4import time5 6# Initialize video capture7cap = cv2.VideoCapture(0)8# Load YOLOv8 model9model = YOLO('last.pt')10 11# Initialize Kalman filter12kalman = cv2.KalmanFilter(4, 2)13kalman.measurementMatrix = np.array([[1, 0, 0, 0], [0, 1, 0, 0]], np.float32)14kalman.transitionMatrix = np.array([[1, 0, 1, 0], [0, 1, 0, 1], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32)15kalman.processNoiseCov = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32) * 0.0316 17# Initialize list to store centroids and the direction line18centroids = []19direction_line_start = None20direction_line_end = None21smoothing_factor = 0.1  # Smoothing factor for direction changes22previous_direction = None23MAX_CENTROIDS = 50  # Maximum number of centroids to keep for drawing24 25while True:26    ret, frame = cap.read()27    if not ret:28        break29 30    results = model(frame, verbose=False)31 32    for result in results:33        if result.masks is not None:34            for mask in result.masks.data:35                mask = mask.cpu().numpy()36 37                if np.any(mask):38                    mask = (mask * 255).astype(np.uint8)39                    mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]))40 41                    M = cv2.moments(mask)42                    if M["m00"] != 0:43                        cX = int(M["m10"] / M["m00"])44                        cY = int(M["m01"] / M["m00"])45                        centroids.append((cX, cY))46 47                        if len(centroids) > MAX_CENTROIDS:48                            centroids.pop(0)49 50                        kalman.correct(np.array([[np.float32(cX)], [np.float32(cY)]]))51                        cv2.circle(frame, (cX, cY), 5, (0, 255, 0), -1)52 53    # Predict the next positions using the Kalman filter54    if len(centroids) > 1:55        for _ in range(10):  # Predict 10 future points56            predicted = kalman.predict()57            pred_x, pred_y = int(predicted[0]), int(predicted[1])58            centroids.append((pred_x, pred_y))59            if len(centroids) > MAX_CENTROIDS:60                centroids.pop(0)61 62    # Update the direction line only if moving backward63    if len(centroids) > 1 and centroids[-1][1] < centroids[-2][1]:64        dx = centroids[-1][0] - centroids[-2][0]65        dy = centroids[-1][1] - centroids[-2][1]66        magnitude = np.sqrt(dx**2 + dy**2)67        if magnitude > 0:68            direction_length = 100  # Length of the direction line69            dir_x = int(dx / magnitude * direction_length)70            dir_y = int(dy / magnitude * direction_length)71            current_direction = (dir_x, dir_y)72            73            if previous_direction is None:74                smoothed_direction = current_direction75            else:76                smoothed_direction = (77                    int(previous_direction[0] * (1 - smoothing_factor) + current_direction[0] * smoothing_factor),78                    int(previous_direction[1] * (1 - smoothing_factor) + current_direction[1] * smoothing_factor)79                )80            81            previous_direction = smoothed_direction82            direction_line_start = (centroids[-1][0], centroids[-1][1])83            direction_line_end = (centroids[-1][0] + smoothed_direction[0], centroids[-1][1] + smoothed_direction[1])84 85    # Draw the direction line86    if direction_line_start and direction_line_end:87        cv2.line(frame, direction_line_start, direction_line_end, (0, 0, 255), 2)88 89    # Show the frame with tracked trajectory and direction lines90    cv2.imshow("Trajectory", frame)91 92    if cv2.waitKey(1) & 0xFF == ord('q'):93        break94 95cap.release()96cv2.destroyAllWindows()