epfl-cs358 / 2024sp-robopong

2 stars 0 forks source link

[CV] Combine paddle and ball tracking into one pipeline #11

Closed Amene-Gafsi closed 5 months ago

Amene-Gafsi commented 5 months ago

This system integrates real-time tracking of both a ball and paddles using the OpenCV library and Python. It employs a color-based tracking method to simultaneously detect and track the largest green-colored circle (representing the ball) and the largest red-colored rectangle (representing the paddles) in video frames. Coordinates for both the ball and paddles are stored and updated continuously.

Algorithm Overview:

  1. Video Capture: Captures video frames from a webcam or video file, defaulting to the webcam if no video path is provided.
  2. Frame Processing: Each frame is resized, blurred, and converted to the HSV color space to standardize input and reduce noise, enhancing the detection of significant objects.
  3. Color Segmentation: Two separate masks are created to isolate green and red hues, which represent the ball and paddles, respectively.
  4. Morphological Operations: Both masks undergo erosion and dilation to clean the image and improve visibility of the tracked objects.
  5. Contour Detection: Contours are extracted from both masks. For the ball, the minimum enclosing circle is computed, and for the paddles, the bounding rectangle method is used to determine dimensions and center.
  6. Data Storage: The centers of the ball and paddles are recorded using deques, which allow for efficient data insertion and retrieval, essential for real-time tracking. To distinguish between the two paddles, an if condition checks if the x-coordinate of the rectangle is in the left part of the frame; if so, it is identified as paddle1, otherwise as paddle2.

Adaptability to Real-World Conditions: The algorithm can be easily adapted for different gaming setups by adjusting parameters : object size and color range.

Code

from collections import deque
from imutils.video import VideoStream
import numpy as np
import argparse
import cv2
import imutils
import time

ap = argparse.ArgumentParser()
ap.add_argument("-v", "--video", help="path to the (optional) video file")
ap.add_argument("-b", "--buffer", type=int, default=64, help="max buffer size")
args = vars(ap.parse_args())

greenLower = (29, 86, 6)
greenUpper = (64, 255, 255)
redLower = (0, 120, 70)
redUpper = (10, 255, 255)

paddle_width_lower = 100
ball_radius_lower = 10
board_center = 300

ball_pts = deque(maxlen=args["buffer"])
paddle1_pts = deque(maxlen=args["buffer"])
paddle2_pts = deque(maxlen=args["buffer"])

if not args.get("video", False):
    vs = VideoStream(src=0).start()
else: vs = cv2.VideoCapture(args["video"])
time.sleep(2.0)

while True:
    frame = vs.read()
    frame = frame[1] if args.get("video", False) else frame
    if frame is None: break

    frame = imutils.resize(frame, width=600) #size of board
    blurred = cv2.GaussianBlur(frame, (11, 11), 0)
    hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)

    ball_mask = cv2.inRange(hsv, greenLower, greenUpper)
    ball_mask = cv2.erode(ball_mask, None, iterations=2)
    ball_mask = cv2.dilate(ball_mask, None, iterations=2)

    paddle_mask = cv2.inRange(hsv, redLower, redUpper)
    paddle_mask = cv2.erode(paddle_mask, None, iterations=2)
    paddle_mask = cv2.dilate(paddle_mask, None, iterations=2)

    ball_cnts = cv2.findContours(ball_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    ball_cnts = imutils.grab_contours(ball_cnts)
    ball_center = None

    paddle_cnts = cv2.findContours(paddle_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    paddle_cnts = imutils.grab_contours(paddle_cnts)
    paddle_center = None

    if len(ball_cnts) > 0:
        ball_c = max(ball_cnts, key=cv2.contourArea)
        ((ball_x, ball_y), ball_radius) = cv2.minEnclosingCircle(ball_c)        
        if ball_radius > ball_radius_lower:
            M = cv2.moments(ball_c)
            ball_center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
            cv2.circle(frame, (int(ball_x), int(ball_y)), int(ball_radius),
                (0, 255, 255), 2)
            cv2.circle(frame, ball_center, 5, (0, 0, 255), -1)

    if len(paddle_cnts) > 0:
        paddle_c = max(paddle_cnts, key=cv2.contourArea)
        paddle_x, paddle_y, paddle_w, paddle_h = cv2.boundingRect(paddle_c)
        if paddle_w > paddle_width_lower:
            paddle_center = (int(paddle_x + paddle_w / 2), int(paddle_y + paddle_h / 2))
            cv2.rectangle(frame, (paddle_x, paddle_y), (paddle_x + paddle_w, paddle_y + paddle_h), (0, 255, 0), 2)
            cv2.circle(frame, paddle_center, 5, (0, 0, 255), -1)

    if ball_center is not None : 
        ball_pts.appendleft(ball_center)
        print("ball: ", ball_center)
    if paddle_center is not None:
        if paddle_center[0] < board_center:
            print("paddle 1: ", paddle_center)
            paddle1_pts.appendleft(paddle_center)
        else:
            print("paddle 2: ", paddle_center)
            paddle2_pts.appendleft(paddle_center)

    cv2.imshow("Frame", frame)
    key = cv2.waitKey(1) & 0xFF
    if key == ord("q"):
        break
if not args.get("video", False):
    vs.stop()
else:
    vs.release()
cv2.destroyAllWindows()