track.py (4258B)
1 #!/usr/bin/python2 2 3 import cv2 4 import math 5 import subprocess 6 import numpy 7 import ghmm 8 import models 9 import simplejson 10 11 path = [] 12 13 14 def diffImg(t0, t1, t2): 15 d1 = cv2.absdiff(t2, t1) 16 d2 = cv2.absdiff(t1, t0) 17 return cv2.bitwise_and(d1, d2) 18 19 def pointer_pos(img): 20 moments = cv2.moments(img) 21 area = moments['m00'] 22 23 if area > 100000: 24 x = moments['m10'] / area 25 y = moments['m01'] / area 26 27 return x, y 28 29 return (None, None) 30 31 def movement_direction(x_delta, y_delta, threshold = 10): 32 if abs(x_delta) > threshold or abs(y_delta) > threshold: 33 degree = math.atan2(y_delta, x_delta) 34 if -0.875 * math.pi <= degree < -0.625 * math.pi: 35 direction = models.UP_RIGHT 36 elif -0.625 * math.pi <= degree < -0.375 * math.pi: 37 direction = models.UP 38 elif -0.375 * math.pi <= degree < -0.125 * math.pi: 39 direction = models.UP_LEFT 40 elif -0.125 * math.pi <= degree < 0.125 * math.pi: 41 direction = models.LEFT 42 elif 0.125 * math.pi <= degree < 0.375 * math.pi: 43 direction = models.DOWN_LEFT 44 elif 0.375 * math.pi <= degree < 0.625 * math.pi: 45 direction = models.DOWN 46 elif 0.625 * math.pi <= degree < 0.875 * math.pi: 47 direction = models.DOWN_RIGHT 48 else: 49 direction = models.RIGHT 50 51 return direction 52 else: 53 return None 54 55 56 def execute(emission_seq, models): 57 considered = [] 58 for model, command in models: 59 res = model.forward(emission_seq) 60 considered.append((res[1][-1], command)) 61 62 max_val, command = max(considered) 63 if max_val >= 0.3: 64 subprocess.call(command) 65 print(command) 66 67 68 def train(emission_seq, model): 69 model.baumWelch(emission_seq) 70 71 def current_image(cam): 72 img = cam.read()[1] 73 img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 74 # Extract by color and make binary image (=black/white) 75 img = cv2.inRange(img, (70, 50, 50), (130, 200, 200)) 76 77 # erode and dilate to reduce noise 78 img = cv2.erode(img, numpy.array([[1] * ed_size] * ed_size), iterations=2) 79 img = cv2.dilate(img, numpy.array([[1] * ed_size] * ed_size), iterations=2) 80 81 return img 82 83 84 # Two variables to determine changes in train-mode 85 train_mode_pre = False 86 train_mode = False 87 train_target = 0 88 89 ed_size = 50 90 91 cam = cv2.VideoCapture(0) 92 93 winName = "Movement Indicator" 94 cv2.namedWindow(winName, cv2.CV_WINDOW_AUTOSIZE) 95 96 img = current_image(cam) 97 98 x1, y1 = pointer_pos(img) 99 100 not_changed = 0 101 102 while True: 103 # if we switched to train mode, delete all prior knowledge 104 if train_mode_pre is False and train_mode is True: 105 path = [] 106 not_changed = 0 107 108 train_mode_pre = train_mode 109 x0 = x1 110 y0 = y1 111 112 img = current_image(cam) 113 x1, y1 = pointer_pos(img) 114 115 if x1 is not None and x0 is not None and y1 is not None and y0 is not None: 116 x_delta = x1 - x0 117 y_delta = y1 - y0 118 119 direction = movement_direction(x_delta, y_delta) 120 if direction is not None: 121 path.append(direction) 122 else: 123 not_changed += 1 124 else: 125 not_changed += 1 126 if not_changed > 5: 127 if len(path) >= 2: 128 print(path) 129 if train_mode == False: 130 execute(ghmm.EmissionSequence(models.sigma, path), models.models) 131 else: 132 models.training_data[train_target].append(path) 133 print("Training model %d" % (train_target,)) 134 train_mode = False 135 print("Leaving training mode") 136 137 with open("".join(('models/', str(train_target), '.train')), 'w') as f: 138 simplejson.dump(models.training_data[train_target], f) 139 140 path = [] 141 not_changed = 0 142 143 cv2.imshow(winName, img) 144 145 key = cv2.waitKey(1) 146 if key == 27: 147 cv2.destroyWindow(winName) 148 break 149 elif key == ord('0'): 150 train_mode = True 151 train_target = 0 152 elif key == ord('1'): 153 train_mode = True 154 train_target = 1 155 elif key == ord('2'): 156 train_mode = True 157 train_target = 2 158 elif key == ord('3'): 159 train_mode = True 160 train_target = 3 161 162 print "Goodbye"