Files
opencv/lasertrack.py
T
2016-11-05 19:10:59 +00:00

255 lines
6.8 KiB
Python
Executable File

# import the necessary packages
from __future__ import print_function
from imutils.video import FPS
import sys
import numpy as np
import argparse
import time
import cv2
import pprint
import math
import struct
import ev3.ev3 as ev3
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from picamera import PiCamera
from piCameraPipeline import PiCameraPipeline
from piCameraPipeline.PiCameraPipeline import RgbProcess
from piCameraPipeline.PiCameraPipeline import RgbProcessorAdapter
from piCameraPipeline.VideoSource import VideoSource
# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-n", "--seconds", type=int, default=0,
help="# of seconds to run")
ap.add_argument("-r", "--framerate", type=int, default=30,
help="Framerate in frame/s")
ap.add_argument("-p", "--with-preview", type=str, default='off',
help="Show camera preview")
ap.add_argument("-e", "--with-ev3", type=str, default='off',
help="Connect to ev3")
ap.add_argument("-c", "--conn-addr", type=str, default='127.0.0.1:5555',
help="Connection dev to ev3")
ap.add_argument("-f", "--filename", type=str, default='piCamera',
help="Choose video source")
ap.add_argument("-x", "--width", type=int, default=320,
help="width")
ap.add_argument("-y", "--height", type=int, default=240,
help="height")
ap.add_argument("-kp", "--kp", type=float, default=6.0, # 1.0
help="kp")
ap.add_argument("-ki", "--ki", type=float, default=2.0, # 0.2
help="ki")
ap.add_argument("-kd", "--kd", type=float, default=2.0, # 8.0
help="kd")
args = vars(ap.parse_args())
videoFile = args["filename"]
framerate = args["framerate"]
seconds = args["seconds"]
with_preview = args["with_preview"] == 'on'
with_ev3 = args["with_ev3"] == 'on'
width = args["width"]
height = args["height"]
connAddr = args["conn_addr"]
kpid = (args["kp"], args["ki"], args["kd"])
class PidControl(object):
def __init__(self, kpid):
self.kpid = kpid
self.errorLast = 0
self.errorAccu = 0
self.pidValue = 0
def paramSet(self, kpid):
self.kpid = kpid
def paramGet(self):
return self.kpid
def reset(self):
self.errorLast = 0
self.errorAccu = 0
def process(self, error):
self.errorAccu += error
d = error - self.errorLast
self.errorLast = error
pid = self.kpid[0]*error + self.kpid[1]*self.errorAccu + self.kpid[2]*d
dpid = pid - self.pidValue
self.pidValue = pid
return (pid, dpid)
class LaserTrack(RgbProcess):
def __init__(self, name, resolution=(320, 240, 3), numEntries=2, next=None, brick=None, kpid=(0,0,0), preview=None):
super(LaserTrack, self).__init__(name, resolution, numEntries, next)
self.motorX = -1000
self.motorY = -1000
self.brick = brick
self.isVisible = False
self.frameSkip = 3
self.frameSkipCounter = 0
self.pidCtrlX = PidControl(kpid)
self.pidCtrlY = PidControl(kpid)
self.preview = preview
def onFirstFrame(self):
self.sendBrick(0, 0)
def onLastFrame(self):
print('onLastFrame')
self.sendBrick(0, 0)
def onProcess(self):
data = self.read(0.1)
if data is not None:
gray = cv2.cvtColor(data,cv2.COLOR_BGR2GRAY)
img1_objects = data.copy()
# Min / max on whole image
minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(gray)
x = maxLoc[0]
y = maxLoc[1]
# get ROI with size roi_size centered at maxLoc
roi_size = 33
roi = self.getROI(gray, x, y, roi_size, roi_size)
# do the blur on ROI
roi_blurred = cv2.medianBlur(roi,5)
# Re-calculate centers on blurred image
minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(roi_blurred)
# adjust centers of unfiltered image (x, y are in global pixel space)
x += maxLoc[0] - int((roi_size-1)/2)
y += maxLoc[1] - int((roi_size-1)/2)
# Calculated dot fall-of over different windows
rect = None
state = 0
ref_color = 1
for roi_size in range(3,67,2):
roi = self.getROI(gray, x, y, roi_size, roi_size)
mean_color = cv2.mean(roi)
if state == 0:
ref_color = mean_color[0]
if ref_color < 128:
break
state = 1
elif state == 1:
km = mean_color[0]/ref_color
if km < 0.7:
rect = [(x-roi_size, y-roi_size),(x+roi_size, y+roi_size)]
break
# Is object found?
errorX = 0
errorY = 0
if rect is not None:
if not self.isVisible:
print ("Visible")
self.isVisible = True
img1_objects = cv2.rectangle(img1_objects,rect[0],rect[1],(255,0,0),1)
errorX = -100.0*(float(x)/self.res[0] - 0.5)
errorY = -100.0*(float(y)/self.res[1] - 0.5)
else:
if self.isVisible:
print ("NOT Visible")
self.isVisible = False
# update PID
(motorX, dmotorX) = self.pidCtrlX.process(errorX)
(motorY, dmotorY) = self.pidCtrlY.process(errorY)
# Update brick
if self.isVisible:
self.sendBrick(dmotorX, dmotorY)
else:
self.sendBrick(0, 0)
# Update preview
showFrame = False
if self.frameSkipCounter > 0:
self.frameSkipCounter -= 1
else:
self.frameSkipCounter = self.frameSkip-1
showFrame = self.preview
if showFrame:
cv2.imshow('Objects', img1_objects)
cv2.waitKey(1)
if self.next is not None:
self.next.write(data)
def sendBrick(self, motorX, motorY):
doSend = motorX != self.motorX
doSend |= motorY != self.motorY
self.motorX = motorX
self.motorY = motorY
if self.brick is not None and doSend:
data = struct.pack("f", motorX)
ev3.system_command.write_mailbox(self.brick, 'MotorX', data)
data = struct.pack("f", motorY)
ev3.system_command.write_mailbox(self.brick, 'MotorY', data)
def getROI(self, img, x, y, w, h):
w2 = int((w-1)/2)
h2 = int((h-1)/2)
xmin = max(0, x-w2)
ymin = max(0, y-h2)
xmax = min(self.res[0]-1, x+w2)
ymax = min(self.res[1]-1, y+h2)
return img[ymin:ymax, xmin:xmax]
if with_ev3:
# Connect to brick
with ev3.EV3(connAddr) as brick:
videoSource = VideoSource(videoFile, (width, height,3), framerate)
proc = LaserTrack('LaserTrackProcess', (width, height, 3), 4, brick=brick, kpid=kpid, preview=with_preview)
videoSource.processorAdd(proc)
videoSource.start()
print("[INFO] sampling THREADED frames from `" + videoFile + "` at " + str(framerate) + " frame/s")
try:
if seconds != 0:
time.sleep(seconds)
else:
while True:
time.sleep(1.0)
except(KeyboardInterrupt):
print ("\nKeyboardInterrupt")
videoSource.stop()
else:
videoSource = VideoSource(videoFile, (width, height,3), framerate)
proc = LaserTrack('LaserTrackProcess', (width, height, 3), 4, kpid=kpid, preview=with_preview)
videoSource.processorAdd(proc)
videoSource.start()
print("[INFO] sampling THREADED frames from `" + videoFile + "` at " + str(framerate) + " frame/s")
try:
if seconds != 0:
time.sleep(seconds)
else:
while True:
time.sleep(1.0)
except(KeyboardInterrupt):
print ("\nKeyboardInterrupt")
videoSource.stop()
# do a bit of cleanup
cv2.destroyAllWindows()
print ("End of program")