CS229-DeepRL/object_detection.py at master · jjw-megha/CS229-DeepRL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
from __future__ import division
import cv2
import numpy as np

class object_detection:

	def __init__(self):
		self.colors = {'man': [200, 72, 72], 'skull': [236,236,236]}
		self.threshold = {'key':0.8, 'door':0.9, 'ladder':0.8}

	def blob_detect(self, img, id):
		mask = np.zeros(img.shape, dtype = "uint8")
		mask[:,:,0] = self.colors[id][0];
		mask[:,:,1] = self.colors[id][1];
		mask[:,:,2] = self.colors[id][2];

		diff = img - mask
		diff[:,:,0] = diff[:,:,1]
		diff[:,:,2] = diff[:,:,1]
		indxs = np.where(diff == 0)
		diff[np.where(diff < 0)] = 0
		diff[np.where(diff > 0)] = 0
		diff[indxs] = 255
		mean_y = np.sum(indxs[0]) / np.shape(indxs[0])[0]
		mean_x = np.sum(indxs[1]) / np.shape(indxs[1])[0]
		return [diff]
		#flipped co-ords due to numpy blob detect

	def template_detect(self, img, id):

		template = cv2.imread('templates/' + id + '.png')

		w = np.shape(template)[1]
		h = np.shape(template)[0]
		res = cv2.matchTemplate(img,template,cv2.TM_CCOEFF_NORMED)
		threshold = self.threshold[id]
		loc = np.where( res >= threshold)
		loc[0].setflags(write=True)
		loc[1].setflags(write=True)
		masks = []
		for i in range(np.shape(loc[0])[0]):
			mask = np.zeros(img.shape, dtype = "uint8")
			a = loc[0][i] ; b = loc[1][i]
			c = loc[0][i] + h; d = loc[1][i] + w
			cv2.rectangle(mask, (b,a), (d,c), (255, 255, 255), -1)
			masks.append(mask)
		return masks


	def detect_objects(self,img):
		object_masks = {}
		object_masks['man'] = self.blob_detect(img,'man')
		object_masks['skull'] = self.blob_detect(img,'skull')
		object_masks['key'] = self.template_detect(img,'key')
		object_masks['door'] = self.template_detect(img,'door')
		object_masks['ladder'] = self.template_detect(img,'ladder')
		return object_masks

	def get_overlap(self,img,goal_mask):
		man_mask = self.blob_detect(img,'man')
		# print(man_mask[0].shape, goal_mask.shape)
		overlap = cv2.bitwise_and(man_mask[0], goal_mask)
		# cv2.imshow('image',overlap)
		# cv2.waitKey(0)
		# cv2.destroyAllWindows()
		bits = np.count_nonzero(overlap)
		total_bits = np.count_nonzero(man_mask)
		# print "Overlap", float(bits)/float(total_bits)
		return float(bits)/float(total_bits) > 0.80

	def to_grayscale(self, img):
		return np.mean(img, axis=2).astype(np.uint8).reshape(img.shape[0],img.shape[1],1)

	def downsample(self, img):
		return img[::2, ::2]

	def preprocess(self, img):
		img = img[30:,:,:]
		return self.to_grayscale(self.downsample(img))

	def get_game_region(self, img):
		return img[30:,:,:]

def main():
	objDet = object_detection()
	img_rgb = cv2.imread('templates/19.png')
	img_score_section = img_rgb[15:20, 55:95, :]
	img_game_section = img_rgb[30:,:,:]
	man_mask = objDet.blob_detect(img_game_section,'man')
	print objDet.get_overlap(img_game_section,man_mask[0])

	# objDet.blob_detect(img_game_section, 'skull')
	# object_masks = objDet.detect_objects(img_game_section)
	# for key in object_masks.keys():
	# 	print key
	# 	for mask in object_masks[key]:
	# 		cv2.imshow('image',mask)
	# 		cv2.waitKey(0)
	# 		cv2.destroyAllWindows()


if __name__ == "__main__":
	main()