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
|
||||
from scipy.spatial import distance as dist, distance
|
||||
|
||||
|
||||
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.team = team
|
||||
self.ID = ID
|
||||
@@ -28,18 +27,31 @@ robotList = []
|
||||
robotMarks = []
|
||||
|
||||
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'
|
||||
if blue < 50 and green > 200 and red > 200:
|
||||
if blue <= 50 and green >= 200 and red >= 200:
|
||||
return 'Yellow'
|
||||
if blue > 200 and green < 50 and red > 200:
|
||||
if blue >= 200 and green <= 50 and red >= 200:
|
||||
return 'Purple'
|
||||
if blue < 50 and green > 220 and red < 50:
|
||||
if blue <= 50 and green >= 220 and red <= 50:
|
||||
return 'Green'
|
||||
if blue <= 50 and green <= 200 and red >= 180:
|
||||
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
|
||||
@@ -48,9 +60,6 @@ def IdentifyCircles(img, circle):
|
||||
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':
|
||||
@@ -131,7 +140,7 @@ def main():
|
||||
assignIDmarks()
|
||||
|
||||
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:
|
||||
print(mark)
|
||||
|
||||
@@ -157,7 +166,3 @@ def main():
|
||||
break
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -73,8 +73,7 @@ def assignIDmarks():
|
||||
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=1,
|
||||
maxRadius=99)
|
||||
circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, 1, minDist=20, param1=50, param2=14, minRadius=20, maxRadius=59)
|
||||
return circles
|
||||
|
||||
|
||||
@@ -100,8 +99,8 @@ def main():
|
||||
global ball
|
||||
|
||||
# Open the video file
|
||||
video_path = "Assets/Images/BotsAndBall.png"
|
||||
cap = cv2.VideoCapture(1)
|
||||
video_path = "Assets/Video/Test2.mp4"
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
|
||||
if not cap.isOpened():
|
||||
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