175 lines
4.6 KiB
Python
175 lines
4.6 KiB
Python
import cv2
|
|
import numpy as np
|
|
from sklearn import cluster
|
|
import os
|
|
# Setting up the blob detector
|
|
params = cv2.SimpleBlobDetector_Params()
|
|
|
|
params.filterByInertia
|
|
params.minInertiaRatio = 0.6
|
|
|
|
detector = cv2.SimpleBlobDetector_create(params)
|
|
|
|
def get_blobs(frame):
|
|
#frame_blurred = cv2.medianBlur(frame, 3)
|
|
#frame_gray = cv2.cvtColor(frame_blurred, cv2.COLOR_BGR2GRAY)
|
|
blobs = detector.detect(frame)
|
|
|
|
return blobs
|
|
|
|
|
|
def get_dice_from_blobs(blobs):
|
|
# Get centroids of all blobs
|
|
X = []
|
|
for b in blobs:
|
|
pos = b.pt
|
|
|
|
if pos != None:
|
|
X.append(pos)
|
|
|
|
X = np.asarray(X)
|
|
|
|
if len(X) > 0:
|
|
# Important to set min_sample to 0, as a dice may only have one dot
|
|
clustering = cluster.DBSCAN(eps=150, min_samples=1).fit(X)
|
|
|
|
# Find the largest label assigned + 1, that's the number of dice found
|
|
num_dice = max(clustering.labels_) + 1
|
|
|
|
dice = []
|
|
|
|
# Calculate centroid of each dice, the average between all a dice's dots
|
|
for i in range(num_dice):
|
|
X_dice = X[clustering.labels_ == i]
|
|
|
|
centroid_dice = np.mean(X_dice, axis=0)
|
|
|
|
dice.append([len(X_dice), *centroid_dice])
|
|
|
|
return dice
|
|
|
|
else:
|
|
return []
|
|
|
|
|
|
|
|
|
|
def overlay_info(frame, dice, blobs):
|
|
# Overlay blobs
|
|
for b in blobs:
|
|
pos = b.pt
|
|
r = b.size / 2
|
|
|
|
cv2.circle(frame, (int(pos[0]), int(pos[1])),
|
|
int(r), (255, 0, 0), 2)
|
|
|
|
# Overlay dice number
|
|
for d in dice:
|
|
# Get textsize for text centering
|
|
textsize = cv2.getTextSize(
|
|
str(d[0]), cv2.FONT_HERSHEY_PLAIN, 3, 2)[0]
|
|
|
|
cv2.putText(frame, str(d[0]),
|
|
(int(d[1] - textsize[0] / 2),
|
|
int(d[2] + textsize[1] / 2)),
|
|
cv2.FONT_HERSHEY_PLAIN, 3, (0, 255, 0), 2)
|
|
|
|
# standard edge detection filter
|
|
def edge_bb_dice(image):
|
|
|
|
gray_img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
blurred = cv2.GaussianBlur(gray_img, (7, 7), 0)
|
|
detected_edges = cv2.Canny(blurred, 150, 200)
|
|
|
|
cropped = gray_img
|
|
filtered = detected_edges
|
|
|
|
nonzero_coords = np.column_stack(np.where(filtered > 0.0))
|
|
|
|
if len(nonzero_coords) > 0:
|
|
min_y, min_x = np.min(nonzero_coords, axis=0)
|
|
max_y, max_x = np.max(nonzero_coords, axis=0)
|
|
|
|
# padding = 20
|
|
# min_x = max(0, min_x - padding // 2)
|
|
# min_y = max(0, min_y - padding // 2)
|
|
# max_x = min(image.shape[1], max_x + padding // 2)
|
|
# max_y = min(image.shape[0], max_y + padding // 2)
|
|
|
|
cropped = gray_img[min_y:max_y, min_x:max_x]
|
|
|
|
#image_with_bb = np.copy(image)
|
|
#cv2.rectangle(image_with_bb, (min_x, min_y), (max_x, max_y), (0, 255, 0), 2)
|
|
|
|
return image, filtered, cropped
|
|
|
|
|
|
|
|
|
|
# Initialize a video feed
|
|
cap = cv2.VideoCapture(1)
|
|
# 16:9 (544x306 -> 30fps with edge_bb_dice)
|
|
cap.set(3, 544)
|
|
cap.set(4, 306)
|
|
|
|
last_dice = 0
|
|
frame_counter = 0
|
|
label_counter = 0
|
|
can_throw = True
|
|
|
|
current_path = f'../bossfight/assets/current_roll.jpg'
|
|
previous_path = f'../bossfight/assets/previous_roll.jpg'
|
|
current_roll = cv2.imread(current_path)
|
|
|
|
while(True):
|
|
# Grab the latest image from the video feed
|
|
ret, frame = cap.read()
|
|
#temp = frame.copy()
|
|
|
|
image, filtered, cropped = edge_bb_dice(frame)
|
|
|
|
# We'll define these later
|
|
blobs = get_blobs(cropped)
|
|
dice = get_dice_from_blobs(blobs)
|
|
out_frame = overlay_info(cropped, dice, blobs)
|
|
|
|
|
|
# cv2.imshow("frame", image)
|
|
# cv2.imshow("frame2", filtered)
|
|
# cv2.imshow("frame3", cropped)
|
|
|
|
for d in dice:
|
|
if last_dice == d[0]:
|
|
frame_counter += 1
|
|
if frame_counter > 17 and can_throw:
|
|
frame_counter = 0
|
|
|
|
cv2.imwrite(current_path, cropped)
|
|
cv2.imwrite(previous_path, current_roll)
|
|
current_roll = cropped
|
|
|
|
label_counter += 1
|
|
s = f'\n{label_counter},{d[0]}'
|
|
with open('../bossfight/dice_roll_data.csv','a') as fd:
|
|
fd.write(s)
|
|
print(s)
|
|
can_throw = False
|
|
else:
|
|
frame_counter = 0
|
|
last_dice = d[0]
|
|
if not dice:
|
|
can_throw = True
|
|
last_dice = 0
|
|
|
|
|
|
|
|
|
|
res = cv2.waitKey(1)
|
|
|
|
# Stop if the user presses "q"
|
|
if res & 0xFF == ord('q'):
|
|
break
|
|
|
|
# When everything is done, release the capture
|
|
cap.release()
|
|
cv2.destroyAllWindows() |