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

width = 320
height = 240

# 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=bool, default=False,
	help="Framerate in frame/s")
ap.add_argument("-f", "--filename", type=str, default='piCamera',
	help="Choose video source")
args = vars(ap.parse_args())

videoFile = args["filename"]
framerate = args["framerate"]
seconds = args["seconds"]
with_ev3 = args["with_ev3"]

class VideoSource(RgbProcessorAdapter):
	def __init__(self, fileName, resolution, framerate):
		super(VideoSource, self).__init__(resolution)
		self.fileName = fileName
		self.resolution = resolution
		self.framerate = framerate
		self.thread = None
		self.camera = None
		self.cancel = False
		if fileName == 'piCamera':
			self.camera = PiCamera()
			self.camera.resolution = (self.resolution[0], self.resolution[1])
			self.camera.framerate = self.framerate
		else:
			self.needRgbConversion = False
			self.thread = Thread(target=self.fileReadThread)

	def fileReadThread(self):
		pass

	def start(self):
		if self.thread is not None:
			self.thread.start()
			pass
		else:
			self.camera.start_recording(self, format="bgr")
		
	def stop(self):
		if self.thread is not None:
			self.cancel = True
			self.thread.join()
		else:
			self.camera.stop_recording()

	def fileReadThread(self):
		cap = cv2.VideoCapture(self.fileName)
		while not self.cancel:
			ret, frame = cap.read()
			if ret:
				self.write(frame)
			time.sleep(1.0/self.framerate)


class MyRgbProcess(RgbProcess):
	def __init__(self, brick, resolution=(320, 240, 3), numEntries=2, next=None):
		super(MyRgbProcess, self).__init__(resolution, numEntries, next)
		self.brick = brick

	def process(self):
		data = self.read(1.0)
		if data is not None:
			gray = cv2.cvtColor(data,cv2.COLOR_BGR2GRAY)

			img1_objects = data.copy()
			if 1:
				# Min / max approach
				gray_blurred = cv2.medianBlur(gray,5)
				minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(gray_blurred)
				x = maxLoc[0]
				y = maxLoc[1]
				if maxVal > 100*minVal:
					img1_objects = cv2.rectangle(img1_objects,(x-5,y-5),(x+5,y+5),(0,0,255),2)
					print(maxLoc)
					if self.brick is not None:
						error = float(x)-160
						data = struct.pack("f", error/2)
						ev3.system_command.write_mailbox(self.brick, 'ErrorX', data)
						error = float(y)-120
						data = struct.pack("f", error/2)
						ev3.system_command.write_mailbox(self.brick, 'ErrorY', data)

			else:
				# Contour approach
				gray_blurred = cv2.GaussianBlur(gray,(5,5),0)
				ret,thresh = cv2.threshold(gray_blurred,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
				# Contours
				(_, contours, hierachy) = cv2.findContours(thresh.copy(),cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
				img1_contours = np.zeros((height,width,3), np.uint8)
				img1_contours = cv2.drawContours(img1_contours, contours, -1, (0,255,0), 1)
				for contour in contours:
					x,y,w,h = cv2.boundingRect(contour)
					roi = gray[y:y+h, x:x+w]
					mean_color = cv2.mean(gray)
					mean_color_roi = cv2.mean(roi)
					if mean_color_roi[0] > 2*mean_color[0]:
						print (mean_color)
						print (mean_color_roi)
						img1_objects = cv2.rectangle(img1_objects,(x,y),(x+w,y+h),(0,0,255),2)

				cv2.imshow('Contours', img1_contours)

			cv2.imshow('Objects', img1_objects)
			cv2.waitKey(1)

			if self.next is not None:
				self.next.write(cv2.cvtColor(gray_blurred,cv2.COLOR_GRAY2BGR))

# Connect to brick

if with_ev3:
	with ev3.EV3() as brick:

		videoSource = VideoSource(videoFile, (width,height,3), framerate)
		proc = MyRgbProcess(brick, (240, 320, 3), 4)
		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(None, (240, 320, 3), 4)
	videoSource.processorAdd(proc)
	videoSource.start()

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

# do a bit of cleanup
cv2.destroyAllWindows()