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 approach
			gray_blurred = cv2.medianBlur(gray,3)
			minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(gray_blurred)
			x = maxLoc[0]
			y = maxLoc[1]
			rect = None
			state = 0
			ref_color = 1
			for roi_size in range(1,15):
				roi = gray[y-roi_size:y+roi_size, x-roi_size:x+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
				
			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
			
			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(cv2.cvtColor(gray_blurred,cv2.COLOR_GRAY2BGR))

	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)

if with_ev3:
	# Connect to brick
	with ev3.EV3() as brick:

		videoSource = VideoSource(videoFile, (width,height,3), framerate)
		proc = MyRgbProcess('MyRgbProcess', (height, width, 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', (height, width, 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()