ACB-CHECKER/Inference.py at main · abelyo252/ACB-CHECKER · 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
"""
Inference : Used for Video Inferencing
Given the amount of given time depleted , this class
try to capture last frame of streaming and perform yolo inference using cv2.DNN module
By: Abel Yohannes
Internship Project for jimma university
"""


import cv2
from time import time
from Camera import Camera
from utils import *
from cvzone.Utils import *
from evaluate import evaluate , rowWriter


class Deep_Vision:
	'''
	This class holds Game information interacting with the Board and Chess Engine
	'''

	def __init__(self):
		self.camera =  Camera()


	def detection_model(self,frame, board):
		# Player 1
		player1 = []
		player2 = []

		start_time = time()

		# Capture frame-by-frame
		#frame = cv2.imread("res/checker.jpg", cv2.IMREAD_COLOR)
		img = frame.copy()
		# do perspective shift, display in 2nd window

		if img.shape[0] > 0 and img.shape[1] > 0:
			# Binarize the photo
			adaptiveThresh, gray = self.camera.clean_Image(img)
			# Black out all pixels outside the border of the chessboard
			mask, approx = self.camera.initialize_mask(adaptiveThresh, img)

			pts = find_outer_corners(img, approx)
			img_orig = do_perspective_transform(mask, pts)

			img_orig = cv2.resize(img_orig, (512, 512), interpolation=cv2.INTER_AREA)
			# save original copy for piece prediction
			img_orig_predict = img_orig.copy()
			imgs = split_chessboard(img_orig)

			#######################################################
			############        YOLO RECOGNITION         ##########
			#######################################################
			test = img_orig_predict.copy()
			# Display the resulting frame
			test = self.camera.format_yolov5(test)
			outs = self.camera.detect(test, self.camera.net)

			class_ids, confidences, boxes = self.camera.wrap_detection(test, outs[0])

			for (classid, confidence, box) in zip(class_ids, confidences, boxes):

				if classid == 0:

					x, y, w, h = box[0], box[1], box[2], box[3]
					cx, cy = int(box[0] + w / 2), int(box[1] + h / 2)
					test = cv2.circle(test, (cx, cy), 5, (0, 0, 255), -1)

					player1.append([cx, cy])

				else:
					x, y, w, h = box[0], box[1], box[2], box[3]
					cx, cy = int(box[0] + w / 2), int(box[1] + h / 2)
					test = cv2.circle(test, (cx, cy), 5, (0, 0, 255), -1)
					player2.append([cx, cy])

				color = self.camera.colors[int(classid) % len(self.camera.colors)]
				cv2.rectangle(test, box, color, 2)
				cv2.rectangle(test, (box[0], box[1] - 20), (box[0] + box[2], box[1]), color, -1)
				cv2.putText(test, self.camera.class_list[classid], (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, .5,
							(0, 0, 0))

			end_time = time()
			sectake = end_time - start_time

			# print(f"Frames Per Second : {fps}")
			start_time = end_time

			# print("player1 detected : ", player1)
			# print("player2 detected : ", player2)

			w, h, _ = test.shape
			dims = list(range(0, w + 1, w // 8))

			for i in dims:
				# Draw Vertical Line
				test = cv2.line(test, (i, 0), (i, w), (255, 0, 0), 2)
				# Draw Horizontal Line
				test = cv2.line(test, (0, i), (w, i), (255, 0, 0), 2)

			if sectake > 0:
				sectake = "Inf Time : %.2f" % sectake
				test, bboxss = putTextRect(test, str(sectake), [10, 20], .5, 1, offset=8, border=3,
										   font=cv2.FONT_HERSHEY_DUPLEX)

			##############################################
			#########       EVALUATE BOARD    ############
			##############################################

			temp_board = evaluate(board, player1, player2)
			return temp_board

		else:
			print("INVALID FRAME DETECTION")
			return 0