computer vision
This commit is contained in:
Mann Patel
@@ -1,160 +0,0 @@
|
|||||||
import cv2
|
|
||||||
import numpy as np
|
|
||||||
from scipy.spatial import distance as dist, distance
|
|
||||||
import main1
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
class Robot:
|
|
||||||
def __init__(self, pos=None, team='-no team!-', ID='-no ID!-'):
|
|
||||||
self.pos = pos if pos is not None else []
|
|
||||||
self.team = team
|
|
||||||
self.ID = ID
|
|
||||||
self.circles = []
|
|
||||||
|
|
||||||
def add_marking(self, circle=None):
|
|
||||||
if circle is None:
|
|
||||||
circle = [0, 0, [0, 0, 0]]
|
|
||||||
self.circles.append(circle)
|
|
||||||
|
|
||||||
|
|
||||||
class Ball:
|
|
||||||
def __init__(self, pos=None):
|
|
||||||
self.pos = pos if pos is not None else []
|
|
||||||
|
|
||||||
|
|
||||||
# Initialize the ball with default position
|
|
||||||
ball = Ball()
|
|
||||||
# Initialize empty lists for robots and ID markings
|
|
||||||
robotList = []
|
|
||||||
robotMarks = []
|
|
||||||
|
|
||||||
def Color_Detection(blue, green, red):
|
|
||||||
if blue > 220 and green < 50 and red < 50:
|
|
||||||
return 'Blue'
|
|
||||||
if blue < 50 and green > 200 and red > 200:
|
|
||||||
return 'Yellow'
|
|
||||||
if blue > 200 and green < 50 and red > 200:
|
|
||||||
return 'Purple'
|
|
||||||
if blue < 50 and green > 220 and red < 50:
|
|
||||||
return 'Green'
|
|
||||||
if blue < 50 and green < 200 and red > 180:
|
|
||||||
return 'Orange'
|
|
||||||
return 'Unidentified'
|
|
||||||
|
|
||||||
|
|
||||||
def IdentifyCircles(img, circle):
|
|
||||||
global ball
|
|
||||||
|
|
||||||
x, y = int(circle[0]), int(circle[1])
|
|
||||||
blue, green, red = img[y, x, 0], img[y, x, 1], img[y, x, 2]
|
|
||||||
color = Color_Detection(blue, green, red)
|
|
||||||
|
|
||||||
# Debugging statements
|
|
||||||
print(f"Circle at ({x}, {y}) with BGR ({blue}, {green}, {red}) detected as {color}")
|
|
||||||
|
|
||||||
if color == 'Blue' or color == 'Yellow':
|
|
||||||
robotList.append(Robot([x, y], color))
|
|
||||||
elif color == 'Green' or color == 'Purple':
|
|
||||||
robotMarks.append([x, y, color])
|
|
||||||
print('ROBOT FOUND')
|
|
||||||
elif color == 'Orange':
|
|
||||||
ball.pos = [x, y]
|
|
||||||
print(f"Ball found at ({x}, {y})")
|
|
||||||
|
|
||||||
|
|
||||||
def assignIDmarks():
|
|
||||||
if robotList is not None:
|
|
||||||
for idx, robot in enumerate(robotList):
|
|
||||||
distances = []
|
|
||||||
|
|
||||||
for i, mark in enumerate(robotMarks):
|
|
||||||
mark_dist = distance.euclidean(mark[:2], robot.pos)
|
|
||||||
distances.append((i, mark_dist))
|
|
||||||
distances.sort(key=lambda x: x[1])
|
|
||||||
closest_marks_indices = [i for i, _ in distances[:4]]
|
|
||||||
robot.circles = [robotMarks[i] for i in closest_marks_indices]
|
|
||||||
robot.ID = idx + 1
|
|
||||||
|
|
||||||
|
|
||||||
def detect_circles(image):
|
|
||||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
||||||
blurred = cv2.GaussianBlur(gray, (9, 9), 0)
|
|
||||||
circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=15,
|
|
||||||
maxRadius=50)
|
|
||||||
return [circles]
|
|
||||||
|
|
||||||
|
|
||||||
def annotate_image(img):
|
|
||||||
for robot in robotList:
|
|
||||||
team_color = "B" if robot.team == 'Blue' else "Y"
|
|
||||||
cv2.putText(img, f'{team_color}', (robot.pos[0] + 20, robot.pos[1] - 40), cv2.FONT_HERSHEY_SIMPLEX, .75,
|
|
||||||
(255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
cv2.putText(img, f'ID{robot.ID}', (robot.pos[0] + 20, robot.pos[1] - 20), cv2.FONT_HERSHEY_SIMPLEX, .75,
|
|
||||||
(255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
cv2.putText(img, f'{robot.pos}', (robot.pos[0] + 20, robot.pos[1]), cv2.FONT_HERSHEY_SIMPLEX, .75,
|
|
||||||
(255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
|
|
||||||
if ball.pos:
|
|
||||||
cv2.putText(img, f'Ball {ball.pos}', (ball.pos[0] + 20, ball.pos[1] + 20), cv2.FONT_HERSHEY_SIMPLEX, 1,
|
|
||||||
(255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
cv2.circle(img, (ball.pos[0], ball.pos[1]), 10, (0, 165, 255), -1) # Orange color for the ball
|
|
||||||
|
|
||||||
|
|
||||||
# Main function
|
|
||||||
def main():
|
|
||||||
global robotList, robotMarks
|
|
||||||
global ball
|
|
||||||
|
|
||||||
while True:
|
|
||||||
# Initialize globals
|
|
||||||
robotList = []
|
|
||||||
robotMarks = []
|
|
||||||
ball = Ball() # Ensure ball is always an instance of Ball
|
|
||||||
|
|
||||||
# Load and process the image
|
|
||||||
imgpath = "Assets/Images/BotsAndBall.png"
|
|
||||||
img = cv2.imread(imgpath)
|
|
||||||
if img is None:
|
|
||||||
print(f"Failed to load image at path: {imgpath}")
|
|
||||||
return
|
|
||||||
|
|
||||||
cv2.imshow("Original Image", img)
|
|
||||||
|
|
||||||
# Detect circles in the image
|
|
||||||
circles = detect_circles(img)
|
|
||||||
|
|
||||||
if circles is not None:
|
|
||||||
circles = np.uint8(np.around(circles))
|
|
||||||
for circle in circles[0, :]:
|
|
||||||
IdentifyCircles(img, circle)
|
|
||||||
cv2.circle(img, (circle[0], circle[1]), circle[2], (0, 255, 0), 2)
|
|
||||||
cv2.circle(img, (circle[0], circle[1]), 2, (0, 0, 255), 3)
|
|
||||||
|
|
||||||
assignIDmarks()
|
|
||||||
|
|
||||||
for robot in robotList:
|
|
||||||
print(f'There is a {robot.team} robot with these ID circles:')
|
|
||||||
for mark in robot.circles:
|
|
||||||
print(mark)
|
|
||||||
|
|
||||||
if ball.pos:
|
|
||||||
print(f'Ball found at {ball.pos}')
|
|
||||||
|
|
||||||
for robot in robotList:
|
|
||||||
if robot.pos:
|
|
||||||
cv2.circle(img, (robot.pos[0], robot.pos[1]), 10, (0, 0, 0), 5)
|
|
||||||
for mark in robot.circles:
|
|
||||||
cv2.circle(img, (mark[0], mark[1]), 10, (0, 0, 0), 5)
|
|
||||||
|
|
||||||
else:
|
|
||||||
print("No circles detected")
|
|
||||||
|
|
||||||
annotate_image(img)
|
|
||||||
cv2.imshow("Annotated Image", img)
|
|
||||||
cv2.waitKey(0)
|
|
||||||
cv2.destroyAllWindows()
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,144 +0,0 @@
|
|||||||
import cv2
|
|
||||||
import numpy as np
|
|
||||||
from scipy.spatial import distance as dist
|
|
||||||
|
|
||||||
# Initialize empty lists for robots and ID markings
|
|
||||||
robotList = []
|
|
||||||
robotMarks = []
|
|
||||||
|
|
||||||
class Robot:
|
|
||||||
def __init__(self, pos=None, team='-no team!-', ID='-no ID!-'):
|
|
||||||
self.pos = pos if pos is not None else []
|
|
||||||
self.team = team
|
|
||||||
self.ID = ID
|
|
||||||
self.circles = [] # ID markings [x, y, color]
|
|
||||||
|
|
||||||
def add_marking(self, circle=None):
|
|
||||||
if circle is None:
|
|
||||||
circle = [0, 0, [0, 0, 0]]
|
|
||||||
self.circles.append(circle)
|
|
||||||
|
|
||||||
class Ball:
|
|
||||||
def __init__(self, pos=None):
|
|
||||||
self.pos = pos if pos is not None else []
|
|
||||||
|
|
||||||
# Initialize the ball with default position
|
|
||||||
ball = Ball()
|
|
||||||
|
|
||||||
def Color_Detection(blue, green, red):
|
|
||||||
if blue > 220 and green < 50 and red < 50:
|
|
||||||
return 'Blue'
|
|
||||||
elif blue < 50 and green > 200 and red > 200:
|
|
||||||
return 'Yellow'
|
|
||||||
elif blue > 200 and green < 50 and red > 200:
|
|
||||||
return 'Purple'
|
|
||||||
elif blue < 50 and green > 220 and red < 50:
|
|
||||||
return 'Green'
|
|
||||||
elif blue < 50 and green < 200 and red > 220:
|
|
||||||
return 'Orange'
|
|
||||||
return f'Unidentified Color R:{red}, G:{green}, B:{blue}'
|
|
||||||
|
|
||||||
def IdentifyCircles(img, circle):
|
|
||||||
global ball
|
|
||||||
|
|
||||||
x, y = int(circle[0]), int(circle[1])
|
|
||||||
blue, green, red = img[y, x, 0], img[y, x, 1], img[y, x, 2]
|
|
||||||
color = Color_Detection(blue, green, red)
|
|
||||||
|
|
||||||
if color == 'Blue' or color == 'Yellow':
|
|
||||||
robotList.append(Robot([x, y], color))
|
|
||||||
elif color == 'Green' or color == 'Purple':
|
|
||||||
robotMarks.append([x, y, color])
|
|
||||||
elif color == 'Orange':
|
|
||||||
ball = Ball([x, y])
|
|
||||||
|
|
||||||
def assignIDmarks():
|
|
||||||
if robotList is not None:
|
|
||||||
for idx, robot in enumerate(robotList):
|
|
||||||
distances = []
|
|
||||||
|
|
||||||
for i, mark in enumerate(robotMarks):
|
|
||||||
mark_dist = dist.euclidean(mark[:2], robot.pos)
|
|
||||||
distances.append((i, mark_dist))
|
|
||||||
distances.sort(key=lambda x: x[1])
|
|
||||||
closest_marks_indices = [i for i, _ in distances[:4]]
|
|
||||||
robot.circles = [robotMarks[i] for i in closest_marks_indices]
|
|
||||||
robot.ID = idx + 1
|
|
||||||
|
|
||||||
def detect_circles(image):
|
|
||||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
||||||
blurred = cv2.GaussianBlur(gray, (9, 9), 0)
|
|
||||||
circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=15, maxRadius=50)
|
|
||||||
return circles
|
|
||||||
|
|
||||||
def annotate_image(img):
|
|
||||||
for robot in robotList:
|
|
||||||
team_color = "B" if robot.team == 'Blue' else "Y"
|
|
||||||
cv2.putText(img, f'{team_color}', (robot.pos[0] + 20, robot.pos[1] - 40), cv2.FONT_HERSHEY_SIMPLEX, .75, (255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
cv2.putText(img, f'\nID{robot.ID}', (robot.pos[0] + 20, robot.pos[1] - 20), cv2.FONT_HERSHEY_SIMPLEX, .75, (255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
cv2.putText(img, f'{robot.pos }', (robot.pos[0] + 20, robot.pos[1]), cv2.FONT_HERSHEY_SIMPLEX, .75, (255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
|
|
||||||
if ball:
|
|
||||||
cv2.putText(img, f'Ball {ball.pos}', (ball.pos[0] + 20, ball.pos[1] + 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
|
|
||||||
# Main function
|
|
||||||
def main():
|
|
||||||
global robotList, robotMarks, ball
|
|
||||||
|
|
||||||
# Initialize globals
|
|
||||||
robotList = []
|
|
||||||
robotMarks = []
|
|
||||||
ball = None
|
|
||||||
|
|
||||||
# Load and process the video
|
|
||||||
video_path = "/Users/mannpatel/Desktop/Robocup/Assets/Video/Test2.mp4"
|
|
||||||
cap = cv2.VideoCapture(video_path)
|
|
||||||
|
|
||||||
while cap.isOpened():
|
|
||||||
ret, frame = cap.read()
|
|
||||||
if not ret:
|
|
||||||
break
|
|
||||||
|
|
||||||
# Reset robot and mark lists for each frame
|
|
||||||
robotList = []
|
|
||||||
robotMarks = []
|
|
||||||
|
|
||||||
# Detect circles in the frame
|
|
||||||
circles = detect_circles(frame)
|
|
||||||
|
|
||||||
if circles is not None:
|
|
||||||
circles = np.uint16(np.around(circles))
|
|
||||||
for circle in circles[0, :]:
|
|
||||||
IdentifyCircles(frame, circle)
|
|
||||||
cv2.circle(frame, (circle[0], circle[1]), circle[2], (0, 255, 0), 2)
|
|
||||||
cv2.circle(frame, (circle[0], circle[1]), 2, (0, 0, 255), 3)
|
|
||||||
|
|
||||||
assignIDmarks()
|
|
||||||
|
|
||||||
# for robot in robotList:
|
|
||||||
# print(f'There is a {robot.team} robot with these ID circles:')
|
|
||||||
# for mark in robot.circles:
|
|
||||||
# print(mark)
|
|
||||||
|
|
||||||
if ball:
|
|
||||||
print(f'Ball found at {ball.pos}')
|
|
||||||
|
|
||||||
for robot in robotList:
|
|
||||||
cv2.circle(frame, (robot.pos[0], robot.pos[1]), 10, (0, 0, 0), 5)
|
|
||||||
for mark in robot.circles:
|
|
||||||
cv2.circle(frame, (mark[0], mark[1]), 10, (0, 0, 0), 5)
|
|
||||||
|
|
||||||
else:
|
|
||||||
print("No circles detected")
|
|
||||||
|
|
||||||
annotate_image(frame)
|
|
||||||
cv2.imshow("Annotated Video", frame)
|
|
||||||
|
|
||||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
||||||
break
|
|
||||||
|
|
||||||
cap.release()
|
|
||||||
cv2.destroyAllWindows()
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,320 +0,0 @@
|
|||||||
# import cv2
|
|
||||||
# import numpy as np
|
|
||||||
# from scipy.spatial import distance as dist, distance
|
|
||||||
#
|
|
||||||
# class Robot:
|
|
||||||
# def __init__(self, pos=None, team='-no team!-', ID='-no ID!-'):
|
|
||||||
# self.pos = pos if pos is not None else []
|
|
||||||
# self.team = team
|
|
||||||
# self.ID = ID
|
|
||||||
# self.circles = []
|
|
||||||
#
|
|
||||||
# def add_marking(self, circle=None):
|
|
||||||
# if circle is None:
|
|
||||||
# circle = [0, 0, [0, 0, 0]]
|
|
||||||
# self.circles.append(circle)
|
|
||||||
#
|
|
||||||
# class Ball:
|
|
||||||
# def __init__(self, pos=None):
|
|
||||||
# self.pos = pos if pos is not None else []
|
|
||||||
#
|
|
||||||
# ball = Ball()
|
|
||||||
# robotList = []
|
|
||||||
# robotMarks = []
|
|
||||||
#
|
|
||||||
# def Color_Detection(blue, green, red):
|
|
||||||
# if blue >= 220 and green <= 50 and red <= 50:
|
|
||||||
# return 'Blue'
|
|
||||||
# if blue <= 50 and green >= 200 and red >= 200:
|
|
||||||
# return 'Yellow'
|
|
||||||
# if blue >= 200 and green <= 50 and red >= 200:
|
|
||||||
# return 'Purple'
|
|
||||||
# if blue <= 50 and green >= 220 and red <= 50:
|
|
||||||
# return 'Green'
|
|
||||||
# if blue <= 50 and green <= 200 and red >= 180:
|
|
||||||
# return 'Orange'
|
|
||||||
# return 'Unidentified'
|
|
||||||
#
|
|
||||||
#
|
|
||||||
# def IdentifyCircles(img, circle):
|
|
||||||
# global ball
|
|
||||||
#
|
|
||||||
# x, y = int(circle[0]), int(circle[1])
|
|
||||||
# blue, green, red = img[y, x, 0], img[y, x, 1], img[y, x, 2]
|
|
||||||
# color = Color_Detection(blue, green, red)
|
|
||||||
#
|
|
||||||
# if color == 'Blue' or color == 'Yellow':
|
|
||||||
# robotList.append(Robot([x, y], color))
|
|
||||||
# elif color == 'Green' or color == 'Purple':
|
|
||||||
# robotMarks.append([x, y, color])
|
|
||||||
# print('ROBOT FOUND')
|
|
||||||
# elif color == 'Orange':
|
|
||||||
# ball.pos = [x, y]
|
|
||||||
# print(f"Ball found at ({x}, {y})")
|
|
||||||
#
|
|
||||||
#
|
|
||||||
# def assignIDmarks():
|
|
||||||
# if robotList is not None:
|
|
||||||
# for idx, robot in enumerate(robotList):
|
|
||||||
# distances = []
|
|
||||||
#
|
|
||||||
# for i, mark in enumerate(robotMarks):
|
|
||||||
# mark_dist = distance.euclidean(mark[:2], robot.pos)
|
|
||||||
# distances.append((i, mark_dist))
|
|
||||||
# distances.sort(key=lambda x: x[1])
|
|
||||||
# closest_marks_indices = [i for i, _ in distances[:4]]
|
|
||||||
# robot.circles = [robotMarks[i] for i in closest_marks_indices]
|
|
||||||
# robot.ID = idx + 1
|
|
||||||
#
|
|
||||||
#
|
|
||||||
# def detect_circles(image):
|
|
||||||
# gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
||||||
# blurred = cv2.GaussianBlur(gray, (9, 9), 0)
|
|
||||||
# circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=15,
|
|
||||||
# maxRadius=50)
|
|
||||||
# return circles
|
|
||||||
#
|
|
||||||
#
|
|
||||||
# def annotate_image(img):
|
|
||||||
# for robot in robotList:
|
|
||||||
# team_color = "B" if robot.team == 'Blue' else "Y"
|
|
||||||
# sting = f'Team: {team_color} | ID: {robot.ID} | POS: {robot.pos}'
|
|
||||||
# cv2.putText(img, sting, (robot.pos[0] + 20, robot.pos[1] - 40), cv2.FONT_HERSHEY_SIMPLEX, .75,
|
|
||||||
# (255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
#
|
|
||||||
# if ball.pos:
|
|
||||||
# cv2.putText(img, f'Ball {ball.pos}', (ball.pos[0] + 20, ball.pos[1] + 20), cv2.FONT_HERSHEY_SIMPLEX, 1,
|
|
||||||
# (255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
#
|
|
||||||
# # Main function
|
|
||||||
# def main():
|
|
||||||
# global robotList, robotMarks
|
|
||||||
# global ball
|
|
||||||
#
|
|
||||||
# while True:
|
|
||||||
# # Initialize globals
|
|
||||||
# robotList = []
|
|
||||||
# robotMarks = []
|
|
||||||
# ball = Ball() # Ensure ball is always an instance of Ball
|
|
||||||
#
|
|
||||||
# # Load and process the image
|
|
||||||
# imgpath = "/Users/mannpatel/Desktop/Project/Computer Vision/Test2.jpeg"
|
|
||||||
# img = cv2.imread(imgpath)
|
|
||||||
# if img is None:
|
|
||||||
# print(f"Failed to load image at path: {imgpath}")
|
|
||||||
# return
|
|
||||||
#
|
|
||||||
# cv2.imshow("Original Image", img)
|
|
||||||
#
|
|
||||||
# # Detect circles in the image
|
|
||||||
# circles = detect_circles(img)
|
|
||||||
#
|
|
||||||
# if circles is not None:
|
|
||||||
# circles = np.uint16(np.around(circles))
|
|
||||||
# for circle in circles[0, :]:
|
|
||||||
# IdentifyCircles(img, circle)
|
|
||||||
# cv2.circle(img, (circle[0], circle[1]), circle[2], (0, 255, 0), 2)
|
|
||||||
# cv2.circle(img, (circle[0], circle[1]), 2, (0, 0, 255), 3)
|
|
||||||
# assignIDmarks()
|
|
||||||
#
|
|
||||||
# for robot in robotList:
|
|
||||||
# print(f'There is a {robot.team} robot with these ID circles:')
|
|
||||||
# for mark in robot.circles:
|
|
||||||
# print(mark)
|
|
||||||
#
|
|
||||||
# if ball.pos:
|
|
||||||
# print(f'Ball found at {ball.pos}')
|
|
||||||
#
|
|
||||||
# for robot in robotList:
|
|
||||||
# if robot.pos:
|
|
||||||
# cv2.circle(img, (robot.pos[0], robot.pos[1]), 10, (0, 0, 0), 5)
|
|
||||||
# for mark in robot.circles:
|
|
||||||
# cv2.circle(img, (mark[0], mark[1]), 10, (0, 0, 0), 5)
|
|
||||||
#
|
|
||||||
# else:
|
|
||||||
# print("No circles detected")
|
|
||||||
#
|
|
||||||
# annotate_image(img)
|
|
||||||
# cv2.imshow("Annotated Image", img)
|
|
||||||
# cv2.waitKey(0)
|
|
||||||
# cv2.destroyAllWindows()
|
|
||||||
#
|
|
||||||
#
|
|
||||||
# if __name__ == "__main__":
|
|
||||||
# main()
|
|
||||||
import cv2
|
|
||||||
import numpy as np
|
|
||||||
from scipy.spatial import distance as dist, distance
|
|
||||||
|
|
||||||
|
|
||||||
class Robot:
|
|
||||||
def __init__(self, pos=None, team='-no team!-', ID='-no ID!-'):
|
|
||||||
self.pos = pos if pos is not None else []
|
|
||||||
self.team = team
|
|
||||||
self.ID = ID
|
|
||||||
self.circles = []
|
|
||||||
|
|
||||||
def add_marking(self, circle=None):
|
|
||||||
if circle is None:
|
|
||||||
circle = [0, 0, [0, 0, 0]]
|
|
||||||
self.circles.append(circle)
|
|
||||||
|
|
||||||
|
|
||||||
class Ball:
|
|
||||||
def __init__(self, pos=None):
|
|
||||||
self.pos = pos if pos is not None else []
|
|
||||||
|
|
||||||
|
|
||||||
# Initialize the ball with default position
|
|
||||||
ball = Ball()
|
|
||||||
# Initialize empty lists for robots and ID markings
|
|
||||||
robotList = []
|
|
||||||
robotMarks = []
|
|
||||||
|
|
||||||
def Color_Detection(blue, green, red):
|
|
||||||
if blue >= 220 and green <= 50 and red <= 50:
|
|
||||||
return 'Blue'
|
|
||||||
if blue <= 50 and green >= 200 and red >= 200:
|
|
||||||
return 'Yellow'
|
|
||||||
if blue >= 200 and green <= 50 and red >= 200:
|
|
||||||
return 'Purple'
|
|
||||||
if blue <= 50 and green >= 220 and red <= 50:
|
|
||||||
return 'Green'
|
|
||||||
if blue <= 50 and green <= 200 and red >= 220:
|
|
||||||
return 'Orange'
|
|
||||||
return 'Unidentified'
|
|
||||||
|
|
||||||
# def Color_Detection(blue, green, red):
|
|
||||||
# if blue == 246 and green == 0 and red == 0:
|
|
||||||
# return 'Blue'
|
|
||||||
# if blue <= 0 and green >= 250 and red > 350:
|
|
||||||
# return 'Yellow'
|
|
||||||
# if blue == 247 and green == 51 and red == 235:
|
|
||||||
# return 'Purple'
|
|
||||||
# if blue == 77 and green == 252 and red == 118:
|
|
||||||
# return 'Green'
|
|
||||||
# if blue == 50 and green == 113 and red == 228:
|
|
||||||
# return 'Orange'
|
|
||||||
# return 'Unidentified'
|
|
||||||
|
|
||||||
|
|
||||||
def IdentifyCircles(img, circle):
|
|
||||||
global ball
|
|
||||||
|
|
||||||
x, y = int(circle[0]), int(circle[1])
|
|
||||||
blue, green, red = img[y, x, 0], img[y, x, 1], img[y, x, 2]
|
|
||||||
color = Color_Detection(blue, green, red)
|
|
||||||
|
|
||||||
# Debugging statements
|
|
||||||
print(f"Circle at ({x}, {y}) with BGR ({blue}, {green}, {red}) detected as {color}")
|
|
||||||
|
|
||||||
if color == 'Blue' or color == 'Yellow':
|
|
||||||
robotList.append(Robot([x, y], color))
|
|
||||||
elif color == 'Green' or color == 'Purple':
|
|
||||||
robotMarks.append([x, y, color])
|
|
||||||
print('ROBOT FOUND')
|
|
||||||
elif color == 'Orange':
|
|
||||||
ball.pos = [x, y]
|
|
||||||
print(f"Ball found at ({x}, {y})")
|
|
||||||
|
|
||||||
|
|
||||||
def assignIDmarks():
|
|
||||||
if robotList is not None:
|
|
||||||
for idx, robot in enumerate(robotList):
|
|
||||||
distances = []
|
|
||||||
|
|
||||||
for i, mark in enumerate(robotMarks):
|
|
||||||
mark_dist = distance.euclidean(mark[:2], robot.pos)
|
|
||||||
distances.append((i, mark_dist))
|
|
||||||
distances.sort(key=lambda x: x[1])
|
|
||||||
closest_marks_indices = [i for i, _ in distances[:4]]
|
|
||||||
robot.circles = [robotMarks[i] for i in closest_marks_indices]
|
|
||||||
robot.ID = idx + 1
|
|
||||||
|
|
||||||
|
|
||||||
def detect_circles(image):
|
|
||||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
||||||
blurred = cv2.GaussianBlur(gray, (9, 9), 0)
|
|
||||||
circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=15,
|
|
||||||
maxRadius=50)
|
|
||||||
return circles
|
|
||||||
|
|
||||||
|
|
||||||
def annotate_image(img):
|
|
||||||
for robot in robotList:
|
|
||||||
team_color = "B" if robot.team == 'Blue' else "Y"
|
|
||||||
cv2.putText(img, f'{team_color}', (robot.pos[0] + 20, robot.pos[1] - 40), cv2.FONT_HERSHEY_SIMPLEX, .75,
|
|
||||||
(255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
cv2.putText(img, f'ID{robot.ID}', (robot.pos[0] + 20, robot.pos[1] - 20), cv2.FONT_HERSHEY_SIMPLEX, .75,
|
|
||||||
(255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
cv2.putText(img, f'{robot.pos}', (robot.pos[0] + 20, robot.pos[1]), cv2.FONT_HERSHEY_SIMPLEX, .75,
|
|
||||||
(255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
|
|
||||||
if ball.pos:
|
|
||||||
cv2.putText(img, f'Ball {ball.pos}', (ball.pos[0] + 20, ball.pos[1] + 20), cv2.FONT_HERSHEY_SIMPLEX, 1,
|
|
||||||
(255, 255, 255), 2, cv2.LINE_AA)
|
|
||||||
cv2.circle(img, (ball.pos[0], ball.pos[1]), 10, (0, 165, 255), -1) # Orange color for the ball
|
|
||||||
|
|
||||||
|
|
||||||
# Main function
|
|
||||||
def main():
|
|
||||||
global robotList, robotMarks
|
|
||||||
global ball
|
|
||||||
|
|
||||||
# Initialize globals
|
|
||||||
robotList = []
|
|
||||||
robotMarks = []
|
|
||||||
ball = Ball() # Ensure ball is always an instance of Ball
|
|
||||||
|
|
||||||
# Load and process the image
|
|
||||||
imgpath = "/Users/mannpatel/Desktop/Robocup/Computer Vision/Template1.png"
|
|
||||||
img = cv2.imread(imgpath)
|
|
||||||
if img is None:
|
|
||||||
print(f"Failed to load image at path: {imgpath}")
|
|
||||||
return
|
|
||||||
|
|
||||||
cv2.imshow("Original Image", img)
|
|
||||||
|
|
||||||
# Detect circles in the image
|
|
||||||
circles = detect_circles(img)
|
|
||||||
|
|
||||||
if circles is not None:
|
|
||||||
circles = np.uint16(np.around(circles))
|
|
||||||
for circle in circles[0, :]:
|
|
||||||
IdentifyCircles(img, circle)
|
|
||||||
cv2.circle(img, (circle[0], circle[1]), circle[2], (0, 255, 0), 2)
|
|
||||||
cv2.circle(img, (circle[0], circle[1]), 2, (0, 0, 255), 3)
|
|
||||||
|
|
||||||
assignIDmarks()
|
|
||||||
|
|
||||||
for robot in robotList:
|
|
||||||
print(f'There is a {robot.team} robot with these ID {robot.ID}')
|
|
||||||
for mark in robot.circles:
|
|
||||||
print(mark)
|
|
||||||
|
|
||||||
if ball.pos:
|
|
||||||
print(f'Ball found at {ball.pos}')
|
|
||||||
|
|
||||||
for robot in robotList:
|
|
||||||
if robot.pos:
|
|
||||||
cv2.circle(img, (robot.pos[0], robot.pos[1]), 10, (0, 0, 0), 5)
|
|
||||||
for mark in robot.circles:
|
|
||||||
cv2.circle(img, (mark[0], mark[1]), 10, (0, 0, 0), 5)
|
|
||||||
|
|
||||||
else:
|
|
||||||
print("No circles detected")
|
|
||||||
|
|
||||||
annotate_image(img)
|
|
||||||
cv2.imshow("Annotated Image", img)
|
|
||||||
|
|
||||||
# Use cv2.waitKey() to display the window until a key is pressed
|
|
||||||
while True:
|
|
||||||
key = cv2.waitKey(1) & 0xFF
|
|
||||||
if key == ord('q'):
|
|
||||||
break
|
|
||||||
|
|
||||||
cv2.destroyAllWindows()
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
BIN
Computer Vision/.DS_Store → ComputerVision/.DS_Store
vendored
BIN
Computer Vision/.DS_Store → ComputerVision/.DS_Store
vendored
Binary file not shown.
@@ -2,9 +2,8 @@ import cv2
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
from scipy.spatial import distance as dist, distance
|
from scipy.spatial import distance as dist, distance
|
||||||
|
|
||||||
|
|
||||||
class Robot:
|
class Robot:
|
||||||
def __init__(self, pos=None, team='-no team!-', ID='-no ID!-'):
|
def __init__(self, pos=None , team=None , ID=None):
|
||||||
self.pos = pos if pos is not None else []
|
self.pos = pos if pos is not None else []
|
||||||
self.team = team
|
self.team = team
|
||||||
self.ID = ID
|
self.ID = ID
|
||||||
@@ -28,18 +27,31 @@ robotList = []
|
|||||||
robotMarks = []
|
robotMarks = []
|
||||||
|
|
||||||
def Color_Detection(blue, green, red):
|
def Color_Detection(blue, green, red):
|
||||||
if blue > 220 and green < 50 and red < 50:
|
if blue >= 220 and green <= 50 and red <= 50:
|
||||||
return 'Blue'
|
return 'Blue'
|
||||||
if blue < 50 and green > 200 and red > 200:
|
if blue <= 50 and green >= 200 and red >= 200:
|
||||||
return 'Yellow'
|
return 'Yellow'
|
||||||
if blue > 200 and green < 50 and red > 200:
|
if blue >= 200 and green <= 50 and red >= 200:
|
||||||
return 'Purple'
|
return 'Purple'
|
||||||
if blue < 50 and green > 220 and red < 50:
|
if blue <= 50 and green >= 220 and red <= 50:
|
||||||
return 'Green'
|
return 'Green'
|
||||||
if blue <= 50 and green <= 200 and red >= 180:
|
if blue <= 50 and green <= 200 and red >= 220:
|
||||||
return 'Orange'
|
return 'Orange'
|
||||||
return 'Unidentified'
|
return 'Unidentified'
|
||||||
|
|
||||||
|
# def Color_Detection(blue, green, red):
|
||||||
|
# if blue == 246 and green == 0 and red == 0:
|
||||||
|
# return 'Blue'
|
||||||
|
# if blue <= 0 and green >= 250 and red > 350:
|
||||||
|
# return 'Yellow'
|
||||||
|
# if blue == 247 and green == 51 and red == 235:
|
||||||
|
# return 'Purple'
|
||||||
|
# if blue == 77 and green == 252 and red == 118:
|
||||||
|
# return 'Green'
|
||||||
|
# if blue == 50 and green == 113 and red == 228:
|
||||||
|
# return 'Orange'
|
||||||
|
# return 'Unidentified'
|
||||||
|
|
||||||
|
|
||||||
def IdentifyCircles(img, circle):
|
def IdentifyCircles(img, circle):
|
||||||
global ball
|
global ball
|
||||||
@@ -48,9 +60,6 @@ def IdentifyCircles(img, circle):
|
|||||||
blue, green, red = img[y, x, 0], img[y, x, 1], img[y, x, 2]
|
blue, green, red = img[y, x, 0], img[y, x, 1], img[y, x, 2]
|
||||||
color = Color_Detection(blue, green, red)
|
color = Color_Detection(blue, green, red)
|
||||||
|
|
||||||
# Debugging statements
|
|
||||||
print(f"Circle at ({x}, {y}) with BGR ({blue}, {green}, {red}) detected as {color}")
|
|
||||||
|
|
||||||
if color == 'Blue' or color == 'Yellow':
|
if color == 'Blue' or color == 'Yellow':
|
||||||
robotList.append(Robot([x, y], color))
|
robotList.append(Robot([x, y], color))
|
||||||
elif color == 'Green' or color == 'Purple':
|
elif color == 'Green' or color == 'Purple':
|
||||||
@@ -131,7 +140,7 @@ def main():
|
|||||||
assignIDmarks()
|
assignIDmarks()
|
||||||
|
|
||||||
for robot in robotList:
|
for robot in robotList:
|
||||||
print(f'There is a {robot.team} robot with these ID circles:')
|
print(f'There is a {robot.team} robot with these ID {robot.ID}')
|
||||||
for mark in robot.circles:
|
for mark in robot.circles:
|
||||||
print(mark)
|
print(mark)
|
||||||
|
|
||||||
@@ -157,7 +166,3 @@ def main():
|
|||||||
break
|
break
|
||||||
|
|
||||||
cv2.destroyAllWindows()
|
cv2.destroyAllWindows()
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -73,8 +73,7 @@ def assignIDmarks():
|
|||||||
def detect_circles(image):
|
def detect_circles(image):
|
||||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||||
blurred = cv2.GaussianBlur(gray, (9, 9), 0)
|
blurred = cv2.GaussianBlur(gray, (9, 9), 0)
|
||||||
circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=1,
|
circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=20, maxRadius=59)
|
||||||
maxRadius=99)
|
|
||||||
return circles
|
return circles
|
||||||
|
|
||||||
|
|
||||||
@@ -100,8 +99,8 @@ def main():
|
|||||||
global ball
|
global ball
|
||||||
|
|
||||||
# Open the video file
|
# Open the video file
|
||||||
video_path = "Assets/Images/BotsAndBall.png"
|
video_path = "Assets/Video/Test2.mp4"
|
||||||
cap = cv2.VideoCapture(1)
|
cap = cv2.VideoCapture(video_path)
|
||||||
|
|
||||||
if not cap.isOpened():
|
if not cap.isOpened():
|
||||||
print(f"Failed to open video file: {video_path}")
|
print(f"Failed to open video file: {video_path}")
|
||||||
BIN
ComputerVision/__pycache__/Image_Processing.cpython-311.pyc
Normal file
BIN
ComputerVision/__pycache__/Image_Processing.cpython-311.pyc
Normal file
Binary file not shown.
BIN
ComputerVision/__pycache__/Video_Processing.cpython-311.pyc
Normal file
BIN
ComputerVision/__pycache__/Video_Processing.cpython-311.pyc
Normal file
Binary file not shown.
BIN
ComputerVision/__pycache__/final.cpython-311.pyc
Normal file
BIN
ComputerVision/__pycache__/final.cpython-311.pyc
Normal file
Binary file not shown.
BIN
ComputerVision/__pycache__/main2.cpython-311.pyc
Normal file
BIN
ComputerVision/__pycache__/main2.cpython-311.pyc
Normal file
Binary file not shown.
BIN
ComputerVision/__pycache__/robocup_ssl_env.cpython-311.pyc
Normal file
BIN
ComputerVision/__pycache__/robocup_ssl_env.cpython-311.pyc
Normal file
Binary file not shown.
298
ComputerVision/robocup_ssl_env.py
Normal file
298
ComputerVision/robocup_ssl_env.py
Normal file
@@ -0,0 +1,298 @@
|
|||||||
|
# import pygame
|
||||||
|
# import numpy as np
|
||||||
|
|
||||||
|
# # Constants based on the image provided
|
||||||
|
# FIELD_WIDTH = 13.4 # meters
|
||||||
|
# FIELD_HEIGHT = 10.4 # meters
|
||||||
|
# GOAL_WIDTH = 1.8 # meters
|
||||||
|
# GOAL_HEIGHT = 1.8 # meters
|
||||||
|
# GOAL_DEPTH = 0.7 # meters
|
||||||
|
# SCALE = 50 # pixels per meter
|
||||||
|
# FPS = 60
|
||||||
|
|
||||||
|
# # Colors
|
||||||
|
# WHITE = (255, 255, 255)
|
||||||
|
# BLACK = (0, 0, 0)
|
||||||
|
# BLUE = (0, 0, 255)
|
||||||
|
# GREEN = (0, 255, 0)
|
||||||
|
# RED = (255, 0, 0)
|
||||||
|
# ORANGE = (255, 165, 0)
|
||||||
|
|
||||||
|
# class RoboCupSSLEnv:
|
||||||
|
# def __init__(self):
|
||||||
|
# pygame.init()
|
||||||
|
# self.screen = pygame.display.set_mode((int(FIELD_WIDTH * SCALE), int(FIELD_HEIGHT * SCALE)))
|
||||||
|
# pygame.display.set_caption("RoboCup SSL Environment")
|
||||||
|
# self.clock = pygame.time.Clock()
|
||||||
|
|
||||||
|
# self.total_reward = 0
|
||||||
|
# self._reset_positions()
|
||||||
|
|
||||||
|
# def _reset_positions(self):
|
||||||
|
# self.robot_pos = np.array([6.7, 5.2])
|
||||||
|
# self.robot_angle = 0
|
||||||
|
# self.ball_pos = np.array([6.7, 3.2])
|
||||||
|
# self.ball_in_possession = False
|
||||||
|
|
||||||
|
# def reset(self):
|
||||||
|
# self.total_reward = 0
|
||||||
|
# self._reset_positions()
|
||||||
|
# return self._get_obs()
|
||||||
|
|
||||||
|
# def _get_obs(self):
|
||||||
|
# return np.array([
|
||||||
|
# self.robot_pos[0], self.robot_pos[1], self.robot_angle,
|
||||||
|
# self.ball_pos[0], self.ball_pos[1], int(self.ball_in_possession)
|
||||||
|
# ])
|
||||||
|
|
||||||
|
# def step(self, action):
|
||||||
|
# if action == 0: # Turn left
|
||||||
|
# self.robot_angle -= np.pi / 18 # Turn 10 degrees
|
||||||
|
# elif action == 1: # Turn right
|
||||||
|
# self.robot_angle += np.pi / 18 # Turn 10 degrees
|
||||||
|
# elif action == 2: # Move forward
|
||||||
|
# self.robot_pos[0] += 0.1 * np.cos(self.robot_angle)
|
||||||
|
# self.robot_pos[1] += 0.1 * np.sin(self.robot_angle)
|
||||||
|
# elif action == 3: # Move backward
|
||||||
|
# self.robot_pos[0] -= 0.1 * np.cos(self.robot_angle)
|
||||||
|
# self.robot_pos[1] -= 0.1 * np.sin(self.robot_angle)
|
||||||
|
# elif action == 4: # Kick
|
||||||
|
# if self.ball_in_possession:
|
||||||
|
# self.ball_pos = self.robot_pos + 2 * np.array([np.cos(self.robot_angle), np.sin(self.robot_angle)])
|
||||||
|
# self.ball_in_possession = False
|
||||||
|
|
||||||
|
# # Ball possession
|
||||||
|
# if not self.ball_in_possession and np.linalg.norm(self.robot_pos - self.ball_pos) < 0.2:
|
||||||
|
# self.ball_in_possession = True
|
||||||
|
|
||||||
|
# # Move ball with robot if in possession
|
||||||
|
# if self.ball_in_possession:
|
||||||
|
# self.ball_pos = self.robot_pos + np.array([0.2 * np.cos(self.robot_angle), 0.2 * np.sin(self.robot_angle)])
|
||||||
|
|
||||||
|
# # Collision with field boundaries
|
||||||
|
# self.robot_pos = np.clip(self.robot_pos, [0, 0], [FIELD_WIDTH, FIELD_HEIGHT])
|
||||||
|
# self.ball_pos = np.clip(self.ball_pos, [0, 0], [FIELD_WIDTH, FIELD_HEIGHT])
|
||||||
|
|
||||||
|
# # Check for goal on the right side
|
||||||
|
# reward = 0
|
||||||
|
# done = False
|
||||||
|
# if self.ball_pos[0] >= FIELD_WIDTH and (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) <= self.ball_pos[1] <= (FIELD_HEIGHT / 2 + GOAL_HEIGHT / 2):
|
||||||
|
# reward += 1 # Scored a goal
|
||||||
|
# self.total_reward += reward
|
||||||
|
# print("---> Goal! Total Reward:", self.total_reward)
|
||||||
|
# self._reset_positions() # Reset player and ball positions
|
||||||
|
|
||||||
|
# return self._get_obs(), reward, done, {}
|
||||||
|
|
||||||
|
# def handle_keys(self):
|
||||||
|
# keys = pygame.key.get_pressed()
|
||||||
|
# if keys[pygame.K_LEFT]:
|
||||||
|
# return 0 # Turn left
|
||||||
|
# elif keys[pygame.K_RIGHT]:
|
||||||
|
# return 1 # Turn right
|
||||||
|
# elif keys[pygame.K_UP]:
|
||||||
|
# return 2 # Move forward
|
||||||
|
# elif keys[pygame.K_DOWN]:
|
||||||
|
# return 3 # Move backward
|
||||||
|
# elif keys[pygame.K_SPACE]:
|
||||||
|
# return 4 # Kick
|
||||||
|
# return -1 # No action
|
||||||
|
|
||||||
|
# def render(self):
|
||||||
|
# self.screen.fill(BLACK) # Clear screen
|
||||||
|
|
||||||
|
# # Draw field
|
||||||
|
# pygame.draw.rect(self.screen, GREEN, pygame.Rect(0, 0, FIELD_WIDTH * SCALE, FIELD_HEIGHT * SCALE))
|
||||||
|
|
||||||
|
# # Draw center line
|
||||||
|
# pygame.draw.line(self.screen, WHITE, (FIELD_WIDTH * SCALE / 2, 0), (FIELD_WIDTH * SCALE / 2, FIELD_HEIGHT * SCALE), 2)
|
||||||
|
|
||||||
|
# # Draw center circle
|
||||||
|
# pygame.draw.circle(self.screen, WHITE, (int(FIELD_WIDTH * SCALE / 2), int(FIELD_HEIGHT * SCALE / 2)), int(1.0 * SCALE), 2)
|
||||||
|
|
||||||
|
# # Draw goals
|
||||||
|
# # Left goal
|
||||||
|
# pygame.draw.rect(self.screen, WHITE, pygame.Rect(0, (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) * SCALE, GOAL_DEPTH * SCALE, GOAL_HEIGHT * SCALE), 2)
|
||||||
|
# pygame.draw.rect(self.screen, WHITE, pygame.Rect(0, (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) * SCALE, GOAL_WIDTH * SCALE, GOAL_HEIGHT * SCALE), 2)
|
||||||
|
|
||||||
|
# # Right goal
|
||||||
|
# pygame.draw.rect(self.screen, RED, pygame.Rect((FIELD_WIDTH - GOAL_DEPTH) * SCALE, (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) * SCALE, GOAL_DEPTH * SCALE, GOAL_HEIGHT * SCALE), 2)
|
||||||
|
# pygame.draw.rect(self.screen, RED, pygame.Rect((FIELD_WIDTH - GOAL_WIDTH) * SCALE, (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) * SCALE, GOAL_WIDTH * SCALE, GOAL_HEIGHT * SCALE), 2)
|
||||||
|
|
||||||
|
# # Draw robot
|
||||||
|
# robot_center = (int(self.robot_pos[0] * SCALE), int(self.robot_pos[1] * SCALE))
|
||||||
|
# pygame.draw.circle(self.screen, BLUE, robot_center, 10)
|
||||||
|
|
||||||
|
# # Draw direction arrow
|
||||||
|
# robot_arrow_end = (robot_center[0] + int(20 * np.cos(self.robot_angle)), robot_center[1] + int(20 * np.sin(self.robot_angle)))
|
||||||
|
# pygame.draw.line(self.screen, BLUE, robot_center, robot_arrow_end, 3)
|
||||||
|
|
||||||
|
# # Draw ball
|
||||||
|
# ball_center = (int(self.ball_pos[0] * SCALE), int(self.ball_pos[1] * SCALE))
|
||||||
|
# pygame.draw.circle(self.screen, ORANGE, ball_center, 8)
|
||||||
|
|
||||||
|
# pygame.display.flip()
|
||||||
|
# self.clock.tick(FPS)
|
||||||
|
|
||||||
|
# def close(self):
|
||||||
|
# pygame.quit()
|
||||||
|
|
||||||
|
# # Usage
|
||||||
|
# env = RoboCupSSLEnv()
|
||||||
|
# obs = env.reset()
|
||||||
|
# done = False
|
||||||
|
|
||||||
|
# while not done:
|
||||||
|
# for event in pygame.event.get():
|
||||||
|
# if event.type == pygame.QUIT:
|
||||||
|
# done = True
|
||||||
|
|
||||||
|
# action = env.handle_keys()
|
||||||
|
# if action != -1:
|
||||||
|
# obs, reward, done, info = env.step(action)
|
||||||
|
# env.render()
|
||||||
|
|
||||||
|
# if env.total_reward >= 25:
|
||||||
|
# done = True
|
||||||
|
|
||||||
|
# env.close()
|
||||||
|
import gym
|
||||||
|
from gym import spaces
|
||||||
|
import pygame
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
# Constants based on the image provided
|
||||||
|
FIELD_WIDTH = 13.4 # meters
|
||||||
|
FIELD_HEIGHT = 10.4 # meters
|
||||||
|
GOAL_WIDTH = 1.8 # meters
|
||||||
|
GOAL_HEIGHT = 1.8 # meters
|
||||||
|
GOAL_DEPTH = 0.7 # meters
|
||||||
|
SCALE = 50 # pixels per meter
|
||||||
|
FPS = 60
|
||||||
|
|
||||||
|
# Colors
|
||||||
|
WHITE = (255, 255, 255)
|
||||||
|
BLACK = (0, 0, 0)
|
||||||
|
BLUE = (0, 0, 255)
|
||||||
|
GREEN = (0, 255, 0)
|
||||||
|
RED = (255, 0, 0)
|
||||||
|
ORANGE = (255, 165, 0)
|
||||||
|
|
||||||
|
class RoboCupSSLEnv(gym.Env):
|
||||||
|
metadata = {'render.modes': ['human']}
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super(RoboCupSSLEnv, self).__init__()
|
||||||
|
|
||||||
|
# Define action and observation space
|
||||||
|
# They must be gym.spaces objects
|
||||||
|
self.action_space = spaces.Discrete(5) # 5 possible actions
|
||||||
|
self.observation_space = spaces.Box(
|
||||||
|
low=np.array([0, 0, -np.pi, 0, 0, 0]),
|
||||||
|
high=np.array([FIELD_WIDTH, FIELD_HEIGHT, np.pi, FIELD_WIDTH, FIELD_HEIGHT, 1]),
|
||||||
|
dtype=np.float32
|
||||||
|
)
|
||||||
|
|
||||||
|
self.total_reward = 0
|
||||||
|
self._reset_positions()
|
||||||
|
|
||||||
|
def _reset_positions(self):
|
||||||
|
self.robot_pos = np.array([6.7, 5.2])
|
||||||
|
self.robot_angle = 0
|
||||||
|
self.ball_pos = np.array([6.7, 3.2])
|
||||||
|
self.ball_in_possession = False
|
||||||
|
|
||||||
|
def reset(self):
|
||||||
|
self.total_reward = 0
|
||||||
|
self._reset_positions()
|
||||||
|
return self._get_obs()
|
||||||
|
|
||||||
|
def _get_obs(self):
|
||||||
|
return np.array([
|
||||||
|
self.robot_pos[0], self.robot_pos[1], self.robot_angle,
|
||||||
|
self.ball_pos[0], self.ball_pos[1], int(self.ball_in_possession)
|
||||||
|
])
|
||||||
|
|
||||||
|
def step(self, action):
|
||||||
|
if action == 0: # Turn left
|
||||||
|
self.robot_angle -= np.pi / 18 # Turn 10 degrees
|
||||||
|
elif action == 1: # Turn right
|
||||||
|
self.robot_angle += np.pi / 18 # Turn 10 degrees
|
||||||
|
elif action == 2: # Move forward
|
||||||
|
self.robot_pos[0] += 0.1 * np.cos(self.robot_angle)
|
||||||
|
self.robot_pos[1] += 0.1 * np.sin(self.robot_angle)
|
||||||
|
elif action == 3: # Move backward
|
||||||
|
self.robot_pos[0] -= 0.1 * np.cos(self.robot_angle)
|
||||||
|
self.robot_pos[1] -= 0.1 * np.sin(self.robot_angle)
|
||||||
|
elif action == 4: # Kick
|
||||||
|
if self.ball_in_possession:
|
||||||
|
self.ball_pos = self.robot_pos + 2 * np.array([np.cos(self.robot_angle), np.sin(self.robot_angle)])
|
||||||
|
self.ball_in_possession = False
|
||||||
|
|
||||||
|
# Ball possession
|
||||||
|
if not self.ball_in_possession and np.linalg.norm(self.robot_pos - self.ball_pos) < 0.2:
|
||||||
|
self.ball_in_possession = True
|
||||||
|
|
||||||
|
# Move ball with robot if in possession
|
||||||
|
if self.ball_in_possession:
|
||||||
|
self.ball_pos = self.robot_pos + np.array([0.2 * np.cos(self.robot_angle), 0.2 * np.sin(self.robot_angle)])
|
||||||
|
|
||||||
|
# Collision with field boundaries
|
||||||
|
self.robot_pos = np.clip(self.robot_pos, [0, 0], [FIELD_WIDTH, FIELD_HEIGHT])
|
||||||
|
self.ball_pos = np.clip(self.ball_pos, [0, 0], [FIELD_WIDTH, FIELD_HEIGHT])
|
||||||
|
|
||||||
|
# Check for goal on the right side
|
||||||
|
reward = 0
|
||||||
|
done = False
|
||||||
|
if self.ball_pos[0] >= FIELD_WIDTH and (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) <= self.ball_pos[1] <= (FIELD_HEIGHT / 2 + GOAL_HEIGHT / 2):
|
||||||
|
reward += 1 # Scored a goal
|
||||||
|
self.total_reward += reward
|
||||||
|
print("---> Goal! Total Reward:", self.total_reward)
|
||||||
|
self._reset_positions() # Reset player and ball positions
|
||||||
|
|
||||||
|
return self._get_obs(), reward, done, {}
|
||||||
|
|
||||||
|
def render(self, mode='human'):
|
||||||
|
if not hasattr(self, 'screen'):
|
||||||
|
pygame.init()
|
||||||
|
self.screen = pygame.display.set_mode((int(FIELD_WIDTH * SCALE), int(FIELD_HEIGHT * SCALE)))
|
||||||
|
pygame.display.set_caption("RoboCup SSL Environment")
|
||||||
|
self.clock = pygame.time.Clock()
|
||||||
|
|
||||||
|
self.screen.fill(BLACK) # Clear screen
|
||||||
|
|
||||||
|
# Draw field
|
||||||
|
pygame.draw.rect(self.screen, GREEN, pygame.Rect(0, 0, FIELD_WIDTH * SCALE, FIELD_HEIGHT * SCALE))
|
||||||
|
|
||||||
|
# Draw center line
|
||||||
|
pygame.draw.line(self.screen, WHITE, (FIELD_WIDTH * SCALE / 2, 0), (FIELD_WIDTH * SCALE / 2, FIELD_HEIGHT * SCALE), 2)
|
||||||
|
|
||||||
|
# Draw center circle
|
||||||
|
pygame.draw.circle(self.screen, WHITE, (int(FIELD_WIDTH * SCALE / 2), int(FIELD_HEIGHT * SCALE / 2)), int(1.0 * SCALE), 2)
|
||||||
|
|
||||||
|
# Draw goals
|
||||||
|
# Left goal
|
||||||
|
pygame.draw.rect(self.screen, WHITE, pygame.Rect(0, (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) * SCALE, GOAL_DEPTH * SCALE, GOAL_HEIGHT * SCALE), 2)
|
||||||
|
pygame.draw.rect(self.screen, WHITE, pygame.Rect(0, (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) * SCALE, GOAL_WIDTH * SCALE, GOAL_HEIGHT * SCALE), 2)
|
||||||
|
|
||||||
|
# Right goal
|
||||||
|
pygame.draw.rect(self.screen, RED, pygame.Rect((FIELD_WIDTH - GOAL_DEPTH) * SCALE, (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) * SCALE, GOAL_DEPTH * SCALE, GOAL_HEIGHT * SCALE), 2)
|
||||||
|
pygame.draw.rect(self.screen, RED, pygame.Rect((FIELD_WIDTH - GOAL_WIDTH) * SCALE, (FIELD_HEIGHT / 2 - GOAL_HEIGHT / 2) * SCALE, GOAL_WIDTH * SCALE, GOAL_HEIGHT * SCALE), 2)
|
||||||
|
|
||||||
|
# Draw robot
|
||||||
|
robot_center = (int(self.robot_pos[0] * SCALE), int(self.robot_pos[1] * SCALE))
|
||||||
|
pygame.draw.circle(self.screen, BLUE, robot_center, 10)
|
||||||
|
|
||||||
|
# Draw direction arrow
|
||||||
|
robot_arrow_end = (robot_center[0] + int(20 * np.cos(self.robot_angle)), robot_center[1] + int(20 * np.sin(self.robot_angle)))
|
||||||
|
pygame.draw.line(self.screen, GREEN, robot_center, robot_arrow_end, 3)
|
||||||
|
|
||||||
|
# Draw ball
|
||||||
|
ball_center = (int(self.ball_pos[0] * SCALE), int(self.ball_pos[1] * SCALE))
|
||||||
|
pygame.draw.circle(self.screen, ORANGE, ball_center, 8)
|
||||||
|
|
||||||
|
pygame.display.flip()
|
||||||
|
self.clock.tick(FPS)
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
pygame.quit()
|
||||||
33
ComputerVision/train_rl_model.py
Normal file
33
ComputerVision/train_rl_model.py
Normal file
@@ -0,0 +1,33 @@
|
|||||||
|
import gymnasium as gym
|
||||||
|
from stable_baselines3 import PPO
|
||||||
|
from stable_baselines3.common.env_checker import check_env
|
||||||
|
from robocup_ssl_env import RoboCupSSLEnv
|
||||||
|
|
||||||
|
# Create environment
|
||||||
|
env = RoboCupSSLEnv()
|
||||||
|
|
||||||
|
# Check if the environment follows the Gym interface
|
||||||
|
check_env(venv)
|
||||||
|
|
||||||
|
# Instantiate the agent
|
||||||
|
model = PPO('MlpPolicy', env, verbose=1)
|
||||||
|
|
||||||
|
# Train the agent
|
||||||
|
model.learn(total_timesteps=10000)
|
||||||
|
|
||||||
|
# Save the model
|
||||||
|
model.save("ppo_robocup_ssl")
|
||||||
|
|
||||||
|
# To reload the trained model
|
||||||
|
# model = PPO.load("ppo_robocup_ssl")
|
||||||
|
|
||||||
|
# Evaluate the trained agent
|
||||||
|
obs = env.reset()
|
||||||
|
for _ in range(1000):
|
||||||
|
action, _states = model.predict(obs, deterministic=True)
|
||||||
|
obs, reward, done, info = env.step(action)
|
||||||
|
env.render()
|
||||||
|
if done:
|
||||||
|
obs = env.reset()
|
||||||
|
|
||||||
|
env.close()
|
||||||
14
main.py
Normal file
14
main.py
Normal file
@@ -0,0 +1,14 @@
|
|||||||
|
import ComputerVision.Image_Processing as file1
|
||||||
|
import ComputerVision.Video_Processing as file2
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
if sys.argv[1] == "file1":
|
||||||
|
if __name__ == '__main__':
|
||||||
|
file1.main()
|
||||||
|
|
||||||
|
elif sys.argv[1] == "file2":
|
||||||
|
if __name__ == '__main__':
|
||||||
|
file2.main()
|
||||||
|
else:
|
||||||
|
print("Invalid argument")
|
||||||
Reference in New Issue
Block a user