PyRobovision
From ICO wiki
import cv2
import numpy as np
import signal
import sys
from flask import Flask, render_template, Response
from collections import deque
from datetime import datetime
from time import time, sleep
from threading import Thread
from PyMata.pymata import PyMata
class Motors(Thread):
MOTOR_1_PWM = 2
MOTOR_1_A = 3
MOTOR_1_B = 4
MOTOR_2_PWM = 5
MOTOR_2_A = 6
MOTOR_2_B = 7
MOTOR_3_PWM = 8
MOTOR_3_A = 9
MOTOR_3_B = 10
def __init__(self):
Thread.__init__(self)
self.daemon = True
self.board = PyMata("/dev/ttyACM1")
def signal_handler(sig, frame):
self.stop_motors()
self.board.reset()
sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)
self.board.set_pin_mode(self.MOTOR_1_PWM, self.board.PWM, self.board.DIGITAL)
self.board.set_pin_mode(self.MOTOR_1_A, self.board.OUTPUT, self.board.DIGITAL)
self.board.set_pin_mode(self.MOTOR_1_B, self.board.OUTPUT, self.board.DIGITAL)
self.board.set_pin_mode(self.MOTOR_2_PWM, self.board.PWM, self.board.DIGITAL)
self.board.set_pin_mode(self.MOTOR_2_A, self.board.OUTPUT, self.board.DIGITAL)
self.board.set_pin_mode(self.MOTOR_2_B, self.board.OUTPUT, self.board.DIGITAL)
self.board.set_pin_mode(self.MOTOR_3_PWM, self.board.PWM, self.board.DIGITAL)
self.board.set_pin_mode(self.MOTOR_3_A, self.board.OUTPUT, self.board.DIGITAL)
self.board.set_pin_mode(self.MOTOR_3_B, self.board.OUTPUT, self.board.DIGITAL)
self.dx, self.dy = 0, 0
def stop_motors(self):
self.board.digital_write(self.MOTOR_1_B, 0)
self.board.digital_write(self.MOTOR_1_A, 0)
self.board.digital_write(self.MOTOR_2_B, 0)
self.board.digital_write(self.MOTOR_2_A, 0)
self.board.digital_write(self.MOTOR_3_B, 0)
self.board.digital_write(self.MOTOR_3_A, 0)
def run(self):
while True:
# Reset all direction pins to avoid damaging H-bridges
self.stop_motors()
dist = abs(self.dx)
if dist > 2:
if self.dx > 0:
print("Turning left")
self.board.digital_write(self.MOTOR_1_B, 1)
self.board.digital_write(self.MOTOR_2_B, 1)
self.board.digital_write(self.MOTOR_3_B, 1)
else:
print("Turning right")
self.board.digital_write(self.MOTOR_1_A, 1)
self.board.digital_write(self.MOTOR_2_A, 1)
self.board.digital_write(self.MOTOR_3_A, 1)
self.board.analog_write(self.MOTOR_1_PWM, int(dist ** 0.7 + 25))
self.board.analog_write(self.MOTOR_2_PWM, int(dist ** 0.7 + 25))
self.board.analog_write(self.MOTOR_3_PWM, int(dist ** 0.7 + 25))
elif self.dy > 30:
print("Going forward")
self.board.digital_write(self.MOTOR_1_B, 1)
self.board.digital_write(self.MOTOR_3_A, 1)
self.board.analog_write(self.MOTOR_1_PWM, int(self.dy ** 0.5 )+30)
self.board.analog_write(self.MOTOR_2_PWM, 0)
self.board.analog_write(self.MOTOR_3_PWM, int(self.dy ** 0.5 )+30)
sleep(0.03)
class FrameGrabber(Thread):
BALL_LOWER = ( 5, 140, 140)
BALL_UPPER = (30, 255, 255)
def __init__(self, width=320, height=240):
Thread.__init__(self)
self.daemon = True
self.video = cv2.VideoCapture(0)
self.video.set(3, width)
self.video.set(4, height)
self.timestamp = time()
self.frames = 0
self.fps = 50
self.current_frame = None
def run(self):
while True:
self.frames += 1
timestamp_begin = time()
if self.frames > 10:
self.fps = self.frames / (timestamp_begin - self.timestamp)
self.frames = 0
self.timestamp = timestamp_begin
success, frame = self.video.read()
frame = cv2.flip(frame, 1)
original = frame
blurred = cv2.blur(frame, (4,4))
hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)
mask = cv2.inRange(hsv, self.BALL_LOWER, self.BALL_UPPER)
mask = cv2.dilate(mask, None, iterations=2)
cutout = cv2.bitwise_and(frame,frame, mask= mask)
cnts = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
if len(cnts) > 0:
c = max(cnts, key=cv2.contourArea)
(x, y), radius = cv2.minEnclosingCircle(c)
M = cv2.moments(c)
center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
if radius > 5:
cv2.circle(frame, (int(x), int(y)), int(radius),
(0, 255, 255), 2)
cv2.circle(frame, center, 5, (0, 0, 255), -1)
radius = 1/radius
radius = round(radius*100*11.35, 2)
cv2.putText(original,str(radius),(int(x),int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.7,(255,255,255),1,cv2.LINE_AA)
cv2.putText(original,str(radius),(int(x+3),int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.59,(0,0,0),1,cv2.LINE_AA)
delta = int((x-160)/5.0)
motors.dx, motors.dy = delta, 240-y
cv2.putText(frame,"%.01f fps" % self.fps, (10,20), cv2.FONT_HERSHEY_SIMPLEX, 0.3,(255,255,255),1,cv2.LINE_AA)
self.current_frame = np.hstack([original, cutout])
motors = Motors()
grabber = FrameGrabber()
motors.start()
grabber.start()
app = Flask(__name__)
@app.route('/')
def index():
return render_template('index.html')
@app.route('/video_feed')
def video_feed():
def generator():
while True:
if grabber.current_frame != None:
ret, jpeg = cv2.imencode('.jpg', grabber.current_frame, (cv2.IMWRITE_JPEG_QUALITY, 10))
yield b'--frame\r\nContent-Type: image/jpeg\r\n\r\n' + jpeg.tobytes() + b'\r\n\r\n'
sleep(0.2)
return Response(generator(),
mimetype='multipart/x-mixed-replace; boundary=frame')
if __name__ == '__main__':
app.run(host='0.0.0.0', debug=True,use_reloader=False,threaded=True)