detect_robots.py

import cv2
import numpy as np
import time
import math

# Read the camera
cap = cv2.VideoCapture(0)

# Define range of colors in HSV
lower = {'Blue':(40 , 170 , 115),'Green':(63 , 50 , 50),'Pink':(125 , 100 , 100)}
upper = {'Blue':(120 , 255 , 210),'Green':(98 , 250 , 140),'Pink':(179 , 255 , 255)}

# Define colors in BGR
colors = {'Blue':(255,0,0),'Green':(0,255,0),'Pink':(255,0,255)}

# Create masks
masks = {}

# Create lists of robots and define color of the ball
positions = {'Blue'}
teams = {'Green','Pink'}

while(1):

	# Take each frame
	_, frame = cap.read()

	# Convert BGR to HSV
	hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
	
	# Define masks
	for key in colors:
		masks[key] = cv2.inRange(hsv, lower[key], upper[key])
	mask_all = np.full((frame.shape[0], frame.shape[1]),0,dtype=np.uint8)
	
	# For each team
	for team in teams:
	
		# For each robot of  each team
		for position in positions:

			# Create a mask for the robot
			mask = cv2.bitwise_or(masks[team], masks[position])

			# Find contours in the mask and initialize the current (x, y) center of the robot
			cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)[-2]
			center = None

			if len(cnts) > 0:

				# Find largest contour in the robot mask
				contour = max(cnts, key = cv2.contourArea)

				# Get circle info
				((x, y), radius) = cv2.minEnclosingCircle(contour)
				M = cv2.moments(contour)
				if M["m00"] != 0:
		 			center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
				else:
					center = None
				
				if center != None:

					# Create a mask for the contour
					mask_c=np.full((frame.shape[0], frame.shape[1]),0,dtype=np.uint8)
					cv2.drawContours(mask_c, [contour], -1, (255,255,255), -1)
					
					# Mask the team color and position color
					mask_team=cv2.bitwise_and(mask_c,mask_c, mask= masks[team])
					mask_position=cv2.bitwise_and(mask_c,mask_c, mask= masks[position])
					
					# Find theirs centers
					cnts_team = cv2.findContours(mask_team.copy(), cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)[-2]
					center_team = None
					if len(cnts_team) > 0:			
						contour_team = max(cnts_team, key = cv2.contourArea)
						((x_team, y_team), radius_team) = cv2.minEnclosingCircle(contour_team)
						M_team = cv2.moments(contour_team)
						if M_team["m00"] != 0:
					 		center_team = (int(M_team["m10"] / M_team["m00"]), int(M_team["m01"] / M_team["m00"]))
						else:
							center_team = None
					else: 
						center_team = None
					cnts_position = cv2.findContours(mask_position.copy(), cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)[-2]
					center_position = None
					if len(cnts_position) > 0:				
						contour_position = max(cnts_position, key = cv2.contourArea)
						((x_position, y_position), radius_position) = cv2.minEnclosingCircle(contour_position)
						M_position = cv2.moments(contour_position)
						if M_position["m00"] != 0:
					 		center_position = (int(M_position["m10"] / M_position["m00"]), int(M_position["m01"] / M_position["m00"]))
						else:
							center_position = None
					else: 
						center_position = None
					if (center_team != None) and (center_position != None):
						# Mask the contour
						cv2.drawContours(mask_all, [contour], -1, (255,255,255), -1)
						
						# Calculate robot's orientation
						angle = math.atan2(center_position[1]-center_team[1],center_position[0]-center_team[0])*180/math.pi
						if angle < 0:
							angle = 360 + angle						
						
						# Print position and angle and save them into info		
						cv2.line(frame, center_position, center_team, colors[team])
						cv2.putText(frame, position, (int(x-radius),int(y-radius)), cv2.FONT_HERSHEY_SIMPLEX, 0.6,colors[team],2)
						cv2.putText(frame, str(center)+', Angle:'+str(int(angle)), (int(x-radius),int(y-radius+20)), cv2.FONT_HERSHEY_SIMPLEX, 0.6,colors[team],2)
						# Bitwise-AND mask and original image
						res = cv2.bitwise_and(frame,frame, mask= mask)
						res_c = cv2.bitwise_and(frame, frame, mask= mask_c)
	
	res_all = cv2.bitwise_and(frame,frame,mask= mask_all)	
	
	# Show screens
	cv2.imshow('frame',frame)
	cv2.imshow('mask_all',mask_all)
	cv2.imshow('res_all',res_all)
	
	k = cv2.waitKey(1) & 0xFF
	if k == ord("q"):
		break
cv2.destroyAllWindows()