ayushkumarshah/airpaint.py

## airpaint.py
import cv2
import numpy as np
from collections import deque

class Airpainter(object):
    """
    An airpainter application.
    The application takes video feed as input and outputs lines tracked.

    Attributes
    ----------
    threshold: tuple
               (lower, upper) contains lower and upper threshold HSV values for
               tracking an object.
    video_cap: cv2.VideoCapture
               VideoCapture constructor.
    kernel: numpy.array
            kernel used for erosion and dilation.
    points: list
            list containing collections.deque object to store detected points.
    source_dim: tuple
                (width, height): width and height of video_cap source.
    paintWindow: numpy.array
                 Paint window for drawing.
    clear_btn: dict
               dictionary containing dimension and color values information for
               clear button and text.
    contour_color: tuple
                   (r, g, b): color value for detected contour.
    """

    def __init__(self, threshold, webcam=False, video_file=None):
        """
        Parameters
        ----------
        threshold: tuple
                   (lower, upper) contains lower and upper threshold HSV values
                   for tracking an object.
        webcam: bool, optional, default: False
                flag to enable or disable webcam as source video.
        video_file: str or None, optional, default: None
                    path to source video file if webcam is set to False.
        """
        self.threshold = threshold
        # Select input video source
        if webcam:
            self.video_cap = cv2.VideoCapture(0)
        else:
            self.video_cap = cv2.VideoCapture(video_file)
        # Define a 5x5 kernel for erosion and dilation
        self.kernel = np.ones((5, 5), np.uint8)
        # Define deque for storing detected points
        self.points = [deque(maxlen=512)]
        # Get width and height of video source
        width  = int(self.video_cap.get(3))
        height = int(self.video_cap.get(4))
        self.source_dim = (width, height)
        # Define a paint window for drawing
        self.paintWindow = np.zeros((self.source_dim[1],
                                     self.source_dim[0], 3)) + 255
        # Define dimensions for clear button and text
        clear_width, clear_height = (100, 64)
        clear_x1, clear_y1 = 90, 10
        clear_x2 = clear_x1 + clear_width
        clear_y2 = clear_y1 + clear_height
        clear_text_x = clear_x1 + 9
        clear_text_y = clear_y1 + int(clear_height / 2)
        # Build dictionary for clear button
        self.clear_btn = {}
        self.clear_btn['dim'] = (clear_width, clear_height)
        self.clear_btn['start'] = (clear_x1, clear_y1)
        self.clear_btn['end'] = (clear_x2, clear_y2)
        self.clear_btn['text'] = {'value': "CLEAR",
                                  'position': (clear_text_x, clear_text_y)}
        # Define colors for clear button, text and circle
        self.clear_btn['bg_color'] = (0, 0, 0) # Black
        self.clear_btn['text_color'] = (255, 255, 255) # White
        self.contour_color = (0, 255, 255) # Yellow
        # Add clear button and text in paint window
        self.add_button()

    def add_button(self, paint=True, frame=None):
        """
        Add clear button in a window

        Parameters
        ----------
        paint: bool, optional, default: True
               flag to select paint window (true) or source window (False).
        frame: numpy.array or None, optional, default: None
               input source frame if paint is set to False.
        """
        if paint:
            self.paintWindow = cv2.rectangle(self.paintWindow.copy(),
                                    self.clear_btn['start'],
                                    self.clear_btn['end'],
                                    self.clear_btn['bg_color'], 2)
            cv2.putText(self.paintWindow, self.clear_btn['text']['value'],
                        self.clear_btn['text']['position'],
                        cv2.FONT_HERSHEY_SIMPLEX, 0.75,
                        self.clear_btn['bg_color'], 2, cv2.LINE_AA)
        else:
            frame =  cv2.rectangle(frame, self.clear_btn['start'],
                                    self.clear_btn['end'],
                                    self.clear_btn['bg_color'], -1)

            cv2.putText(frame, self.clear_btn['text']['value'],
                        self.clear_btn['text']['position'],
                        cv2.FONT_HERSHEY_SIMPLEX, 0.75,
                        self.clear_btn['text_color'], 2, cv2.LINE_AA)

    def init_op_vid(self):
        """
        Initialize output video writers.

        Returns
        -------
        output_source: cv2.VideoWriter
                       Video Writer to save tracking output frames in source window.
        output_paint: cv2.VideoWriter
                      Video Writer to save tracking output frames in paint window.
        """
        # Change codec according to the system
        fourcc = cv2.VideoWriter_fourcc('X','V','I','D')
        output_source = cv2.VideoWriter('output_source.mkv',
                                fourcc, 20, self.source_dim)
        output_paint = cv2.VideoWriter('output_paint.mkv',
                                fourcc, 20, self.source_dim)
        return output_source, output_paint


    def find_object(self, hsv):
        """
        Return pixels within defined thresholds.

        Parameters
        ----------
        hsv: numpy.ndarray
             input source frame in hsv color space.

        Returns
        -------
        mask: numpy.array
              mask representing pixels between defined thresholds.
        """
        # Determine pixels falling within the defined thresholds
        mask = cv2.inRange(hsv, self.threshold[0], self.threshold[1],)
        # Blur and dilate the binary image for clear view
        mask = cv2.erode(mask, self.kernel, iterations=2)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, self.kernel)
        mask = cv2.dilate(mask, self.kernel, iterations=1)
        return mask

    def draw_contour(self, frame, contours):
        """
        Draw and return the largest contour among detected contours in the
        source frame.

        Parameters
        ----------
        frame: numpy.ndarray
               source window frame.
        contours: list
                  list of detected contours.

        Returns
        -------
        contour: numpy.array
                 largest contour.
        """
        # Sort the contours and find the largest one
        contour = sorted(contours, key = cv2.contourArea, reverse = True)[0]
        # Get the radius of the enclosing circle around the contour
        ((x, y), radius) = cv2.minEnclosingCircle(contour)
        # Draw the circle around the contour
        cv2.circle(frame, (int(x), int(y)),
                    int(radius), self.contour_color, 2)
        return contour

    def draw_lines(self, frame, paintWindow):
        """
        Draw lines through the tracked path in source and paint window.

        Parameters
        ----------
        frame: numpy.ndarray
               source window frame.
        paintWindow: numpy.ndarray
                     paint window frame.

        """
        for index in range(len(self.points)):
            for j in range(1, len(self.points[index])):
                if self.points[index][j-1] and self.points[index][j]:
                    cv2.line(frame, self.points[index][j-1],
                                self.points[index][j], color=(0, 0, 255),
                                thickness=2)
                    cv2.line(paintWindow, self.points[index][j - 1],
                                self.points[index][j], color= (0, 0, 255),
                                thickness=2)

    def is_clear(self, center):
        """
        Check if clear button is activated.

        Parameters
        ----------
        center: tuple
                coordinates of the center of contour.

        Returns
        -------
        check: bool
               True if clear button activated, else False.

        """
        check = (self.clear_btn['start'][0] <= center[0] <= self.clear_btn['end'][0]
                    and self.clear_btn['start'][1] <= center[1] <= self.clear_btn['end'][1])
        return check

    def paint(self, save_video=True):
        """
        Paint lines following the tracked object in source and paint window.

        Parameters
        ----------
        save_video: bool, optional, default: True
                    flag to enable or disable writing output video to disk.
        """

        # Initialize paintwindow and index
        paintWindow = self.paintWindow.copy()
        index = 0
        # Create Video Writers to save tracking outputs
        if save_video:
            output_source, output_paint = self.init_op_vid()
        while self.video_cap.isOpened():
            # Capture frame-by-frame
            ret, frame = self.video_cap.read()
            # Exit loop if end of the video reached
            if not ret:
                break
            # Flip left and right side to avoid mirroring
            frame = cv2.flip(frame, 1)
            # Convert BGR to HSV color space
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            # Add Clear button and text in input video screen
            self.add_button(paint=False, frame=frame)
            # Find object using threshold values
            mask = self.find_object(hsv)
            # Find contours in the image
            (contours, _) = cv2.findContours(mask.copy(),
                                            cv2.RETR_EXTERNAL,
                                            cv2.CHAIN_APPROX_SIMPLE)
            center = None
            # Check if contours are found
            if len(contours) > 0:
                contour = self.draw_contour(frame, contours)
                # Get the moments to calculate the center of the contour
                M = cv2.moments(contour)
                center = (int(M['m10'] / M['m00']), int(M['m01'] / M['m00']))
                # Define task of Clear button activation
                if (self.is_clear(center)):
                    self.points = [deque(maxlen=512)]
                    paintWindow = self.paintWindow.copy()
                    index = 0
                else:
                    self.points[index].appendleft(center)
            # Reset if no contours found
            else:
                self.points.append(deque(maxlen=512))
                index += 1
            # Draw lines using tracked points
            self.draw_lines(frame, paintWindow)
            ## Display the image
            cv2.imshow('Tracking', frame)
            cv2.imshow("Paint", paintWindow)
            if save_video:
                paintWindow = paintWindow.astype('uint8')
                output_paint.write(paintWindow)
                output_source.write(frame)
            # Press q to quit
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        # Release the videos
        self.video_cap.release()
        if save_video:
            output_source.release()
            output_paint.release()
        cv2.destroyAllWindows()

def main():
    # Define thresholds
    lower = np.array([90, 80, 30])
    upper = np.array([150, 255, 255])
    threshold = (lower, upper)
    # Define airpainter object
    airpainter = Airpainter(threshold, webcam=True)
    airpainter.paint()

if __name__ == '__main__':
    main()
	import cv2
	import numpy as np
	from collections import deque

	class Airpainter(object):
	"""
	An airpainter application.
	The application takes video feed as input and outputs lines tracked.

	Attributes
	----------
	threshold: tuple
	(lower, upper) contains lower and upper threshold HSV values for
	tracking an object.
	video_cap: cv2.VideoCapture
	VideoCapture constructor.
	kernel: numpy.array
	kernel used for erosion and dilation.
	points: list
	list containing collections.deque object to store detected points.
	source_dim: tuple
	(width, height): width and height of video_cap source.
	paintWindow: numpy.array
	Paint window for drawing.
	clear_btn: dict
	dictionary containing dimension and color values information for
	clear button and text.
	contour_color: tuple
	(r, g, b): color value for detected contour.
	"""

	def __init__(self, threshold, webcam=False, video_file=None):
	"""
	Parameters
	----------
	threshold: tuple
	(lower, upper) contains lower and upper threshold HSV values
	for tracking an object.
	webcam: bool, optional, default: False
	flag to enable or disable webcam as source video.
	video_file: str or None, optional, default: None
	path to source video file if webcam is set to False.
	"""
	self.threshold = threshold
	# Select input video source
	if webcam:
	self.video_cap = cv2.VideoCapture(0)
	else:
	self.video_cap = cv2.VideoCapture(video_file)
	# Define a 5x5 kernel for erosion and dilation
	self.kernel = np.ones((5, 5), np.uint8)
	# Define deque for storing detected points
	self.points = [deque(maxlen=512)]
	# Get width and height of video source
	width = int(self.video_cap.get(3))
	height = int(self.video_cap.get(4))
	self.source_dim = (width, height)
	# Define a paint window for drawing
	self.paintWindow = np.zeros((self.source_dim[1],
	self.source_dim[0], 3)) + 255
	# Define dimensions for clear button and text
	clear_width, clear_height = (100, 64)
	clear_x1, clear_y1 = 90, 10
	clear_x2 = clear_x1 + clear_width
	clear_y2 = clear_y1 + clear_height
	clear_text_x = clear_x1 + 9
	clear_text_y = clear_y1 + int(clear_height / 2)
	# Build dictionary for clear button
	self.clear_btn = {}
	self.clear_btn['dim'] = (clear_width, clear_height)
	self.clear_btn['start'] = (clear_x1, clear_y1)
	self.clear_btn['end'] = (clear_x2, clear_y2)
	self.clear_btn['text'] = {'value': "CLEAR",
	'position': (clear_text_x, clear_text_y)}
	# Define colors for clear button, text and circle
	self.clear_btn['bg_color'] = (0, 0, 0) # Black
	self.clear_btn['text_color'] = (255, 255, 255) # White
	self.contour_color = (0, 255, 255) # Yellow
	# Add clear button and text in paint window
	self.add_button()

	def add_button(self, paint=True, frame=None):
	"""
	Add clear button in a window

	Parameters
	----------
	paint: bool, optional, default: True
	flag to select paint window (true) or source window (False).
	frame: numpy.array or None, optional, default: None
	input source frame if paint is set to False.
	"""
	if paint:
	self.paintWindow = cv2.rectangle(self.paintWindow.copy(),
	self.clear_btn['start'],
	self.clear_btn['end'],
	self.clear_btn['bg_color'], 2)
	cv2.putText(self.paintWindow, self.clear_btn['text']['value'],
	self.clear_btn['text']['position'],
	cv2.FONT_HERSHEY_SIMPLEX, 0.75,
	self.clear_btn['bg_color'], 2, cv2.LINE_AA)
	else:
	frame = cv2.rectangle(frame, self.clear_btn['start'],
	self.clear_btn['end'],
	self.clear_btn['bg_color'], -1)

	cv2.putText(frame, self.clear_btn['text']['value'],
	self.clear_btn['text']['position'],
	cv2.FONT_HERSHEY_SIMPLEX, 0.75,
	self.clear_btn['text_color'], 2, cv2.LINE_AA)

	def init_op_vid(self):
	"""
	Initialize output video writers.

	Returns
	-------
	output_source: cv2.VideoWriter
	Video Writer to save tracking output frames in source window.
	output_paint: cv2.VideoWriter
	Video Writer to save tracking output frames in paint window.
	"""
	# Change codec according to the system
	fourcc = cv2.VideoWriter_fourcc('X','V','I','D')
	output_source = cv2.VideoWriter('output_source.mkv',
	fourcc, 20, self.source_dim)
	output_paint = cv2.VideoWriter('output_paint.mkv',
	fourcc, 20, self.source_dim)
	return output_source, output_paint


	def find_object(self, hsv):
	"""
	Return pixels within defined thresholds.

	Parameters
	----------
	hsv: numpy.ndarray
	input source frame in hsv color space.

	Returns
	-------
	mask: numpy.array
	mask representing pixels between defined thresholds.
	"""
	# Determine pixels falling within the defined thresholds
	mask = cv2.inRange(hsv, self.threshold[0], self.threshold[1],)
	# Blur and dilate the binary image for clear view
	mask = cv2.erode(mask, self.kernel, iterations=2)
	mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, self.kernel)
	mask = cv2.dilate(mask, self.kernel, iterations=1)
	return mask

	def draw_contour(self, frame, contours):
	"""
	Draw and return the largest contour among detected contours in the
	source frame.

	Parameters
	----------
	frame: numpy.ndarray
	source window frame.
	contours: list
	list of detected contours.

	Returns
	-------
	contour: numpy.array
	largest contour.
	"""
	# Sort the contours and find the largest one
	contour = sorted(contours, key = cv2.contourArea, reverse = True)[0]
	# Get the radius of the enclosing circle around the contour
	((x, y), radius) = cv2.minEnclosingCircle(contour)
	# Draw the circle around the contour
	cv2.circle(frame, (int(x), int(y)),
	int(radius), self.contour_color, 2)
	return contour

	def draw_lines(self, frame, paintWindow):
	"""
	Draw lines through the tracked path in source and paint window.

	Parameters
	----------
	frame: numpy.ndarray
	source window frame.
	paintWindow: numpy.ndarray
	paint window frame.

	"""
	for index in range(len(self.points)):
	for j in range(1, len(self.points[index])):
	if self.points[index][j-1] and self.points[index][j]:
	cv2.line(frame, self.points[index][j-1],
	self.points[index][j], color=(0, 0, 255),
	thickness=2)
	cv2.line(paintWindow, self.points[index][j - 1],
	self.points[index][j], color= (0, 0, 255),
	thickness=2)

	def is_clear(self, center):
	"""
	Check if clear button is activated.

	Parameters
	----------
	center: tuple
	coordinates of the center of contour.

	Returns
	-------
	check: bool
	True if clear button activated, else False.

	"""
	check = (self.clear_btn['start'][0] <= center[0] <= self.clear_btn['end'][0]
	and self.clear_btn['start'][1] <= center[1] <= self.clear_btn['end'][1])
	return check

	def paint(self, save_video=True):
	"""
	Paint lines following the tracked object in source and paint window.

	Parameters
	----------
	save_video: bool, optional, default: True
	flag to enable or disable writing output video to disk.
	"""

	# Initialize paintwindow and index
	paintWindow = self.paintWindow.copy()
	index = 0
	# Create Video Writers to save tracking outputs
	if save_video:
	output_source, output_paint = self.init_op_vid()
	while self.video_cap.isOpened():
	# Capture frame-by-frame
	ret, frame = self.video_cap.read()
	# Exit loop if end of the video reached
	if not ret:
	break
	# Flip left and right side to avoid mirroring
	frame = cv2.flip(frame, 1)
	# Convert BGR to HSV color space
	hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
	# Add Clear button and text in input video screen
	self.add_button(paint=False, frame=frame)
	# Find object using threshold values
	mask = self.find_object(hsv)
	# Find contours in the image
	(contours, _) = cv2.findContours(mask.copy(),
	cv2.RETR_EXTERNAL,
	cv2.CHAIN_APPROX_SIMPLE)
	center = None
	# Check if contours are found
	if len(contours) > 0:
	contour = self.draw_contour(frame, contours)
	# Get the moments to calculate the center of the contour
	M = cv2.moments(contour)
	center = (int(M['m10'] / M['m00']), int(M['m01'] / M['m00']))
	# Define task of Clear button activation
	if (self.is_clear(center)):
	self.points = [deque(maxlen=512)]
	paintWindow = self.paintWindow.copy()
	index = 0
	else:
	self.points[index].appendleft(center)
	# Reset if no contours found
	else:
	self.points.append(deque(maxlen=512))
	index += 1
	# Draw lines using tracked points
	self.draw_lines(frame, paintWindow)
	## Display the image
	cv2.imshow('Tracking', frame)
	cv2.imshow("Paint", paintWindow)
	if save_video:
	paintWindow = paintWindow.astype('uint8')
	output_paint.write(paintWindow)
	output_source.write(frame)
	# Press q to quit
	if cv2.waitKey(1) & 0xFF == ord('q'):
	break
	# Release the videos
	self.video_cap.release()
	if save_video:
	output_source.release()
	output_paint.release()
	cv2.destroyAllWindows()

	def main():
	# Define thresholds
	lower = np.array([90, 80, 30])
	upper = np.array([150, 255, 255])
	threshold = (lower, upper)
	# Define airpainter object
	airpainter = Airpainter(threshold, webcam=True)
	airpainter.paint()

	if __name__ == '__main__':
	main()