opencv/lasertrack.py

# 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 threading import Thread

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=10,
	help="# of seconds to run")
ap.add_argument("-r", "--framerate", type=int, default=30,
	help="Framerate in frame/s")
ap.add_argument("-e", "--with-ev3", type=str, default='off',
	help="Framerate in frame/s")
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")
args = vars(ap.parse_args())

videoFile = args["filename"]
framerate = args["framerate"]
seconds = args["seconds"]
with_ev3 = args["with_ev3"] == 'on'
width = args["width"]
height = args["height"]

class MyRgbProcess(RgbProcess):
	def __init__(self, name, resolution=(320, 240, 3), numEntries=2, next=None, brick=None):
		super(MyRgbProcess, self).__init__(name, resolution, numEntries, next)
		self.brick = brick
		self.isVisible = False
		self.frameSkip = 3
		self.frameSkipCounter = 0
	def onFirstFrame(self):
		self.sendBrick(0, 0)

	def onLastFrame(self):
		self.sendBrick(0, 0)

	def onProcess(self):
		data = self.read(1.0)
		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 = 21
			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,33,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?
			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 = float(x)-self.res[0]/2
				errorY = float(y)-self.res[1]/2
				self.sendBrick(errorX/2, errorY/2)
			else:
				if self.isVisible:
					print ("NOT Visible")
					self.sendBrick(0, 0)
				self.isVisible = False
			
			# Update preview
			showFrame = False
			if self.frameSkipCounter > 0:
				self.frameSkipCounter -= 1
			else:
				self.frameSkipCounter = self.frameSkip-1
				showFrame = True

			if showFrame:
				cv2.imshow('Objects', img1_objects)
				cv2.waitKey(1)

			if self.next is not None:
				self.next.write(data)

	def sendBrick(self, errorX, errorY):
		if self.brick is not None:
			data = struct.pack("f", errorX)
			ev3.system_command.write_mailbox(self.brick, 'ErrorX', data)
			data = struct.pack("f", errorY)
			ev3.system_command.write_mailbox(self.brick, 'ErrorY', 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() as brick:

		videoSource = VideoSource(videoFile, (width, height,3), framerate)
		proc = MyRgbProcess('MyRgbProcess', (width, height, 3), 4, brick=brick)
		videoSource.processorAdd(proc)
		videoSource.start()

		print("[INFO] sampling THREADED frames from `" + videoFile + "` at " + str(framerate) + " frame/s")
		time.sleep(seconds)
		videoSource.stop()
		proc.stop()
else:
	videoSource = VideoSource(videoFile, (width, height,3), framerate)
	proc = MyRgbProcess('MyRgbProcess', (width, height, 3), 4)
	videoSource.processorAdd(proc)
	videoSource.start()

	print("[INFO] sampling THREADED frames from `" + videoFile + "` at " + str(framerate) + " frame/s")
	time.sleep(seconds)
	videoSource.stop()

# do a bit of cleanup
cv2.destroyAllWindows()