-
Notifications
You must be signed in to change notification settings - Fork 1
/
pupil.py
54 lines (43 loc) · 1.74 KB
/
pupil.py
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
import numpy as np
import cv2
class Pupil(object):
"""
This class detects the iris of an eye and estimates
the position of the pupil
"""
def __init__(self, eye_frame, threshold):
self.iris_frame = None
self.threshold = threshold
self.x = None
self.y = None
self.detect_iris(eye_frame)
@staticmethod
def image_processing(eye_frame, threshold):
"""Performs operations on the eye frame to isolate the iris
Arguments:
eye_frame (numpy.ndarray): Frame containing an eye and nothing else
threshold (int): Threshold value used to binarize the eye frame
Returns:
A frame with a single element representing the iris
"""
kernel = np.ones((3, 3), np.uint8)
new_frame = cv2.bilateralFilter(eye_frame, 10, 15, 15)
new_frame = cv2.erode(new_frame, kernel, iterations=3)
new_frame = cv2.threshold(new_frame, threshold, 255, cv2.THRESH_BINARY)[1]
return new_frame
def detect_iris(self, eye_frame):
"""Detects the iris and estimates the position of the iris by
calculating the centroid.
Arguments:
eye_frame (numpy.ndarray): Frame containing an eye and nothing else
"""
self.iris_frame = self.image_processing(eye_frame, self.threshold)
contours, _ = cv2.findContours(self.iris_frame, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)[-2:]
contours = sorted(contours, key=cv2.contourArea)
#print(contours)
try:
moments = cv2.moments(contours[-2])
self.x = int(moments['m10'] / moments['m00'])
self.y = int(moments['m01'] / moments['m00'])
except (IndexError, ZeroDivisionError):
pass