main_mini.py

# Leechy Prototype Spectrum Analyzer. 
# Important: MAKE SURE KEYBOARD IS ON ENGLISH AND CAPSLOCK IS NOT ON!

import cv2,pickle,xlsxwriter,time,datetime,os, os.path
from imutils import rotate_bound
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button
from PIL import Image, ImageTk

print('\033[1m' + '\033[91m' + '\nLeechy Labs | Spectrum Analyser\n' + '\033[0m')

drunkdatanum = 0
not_drunkdatanum = 0

DRUNK_DIR = 'data/graph/drunk'
drunkdatanum = len([name for name in os.listdir(DRUNK_DIR) if os.path.isfile(os.path.join(DRUNK_DIR, name))])

NOT_DRUNK_DIR = 'data/graph/not_drunk'
not_drunkdatanum = len([name for name in os.listdir(DRUNK_DIR) if os.path.isfile(os.path.join(NOT_DRUNK_DIR, name))])

row = 1

sensetivity = 615
x1 = 480
y1 = 1000
x2 = 578
y2 = 1040
angle = 52.5
flip = False
plot_graph = False

cmp_graph = [500]
nor_color = (0, 0, 1)
mouse_rectangle = False
clear = True
live = False
auto = False

plt.ion()

def rotate_clockwise(matrix, degree=90):
	try: return matrix if not degree else rotate_clockwise(zip(*matrix[::-1]), degree-90)
	except TypeError: return [[0]]

def mouse_event(event,x,y,flags,param):
	global x1, x2, y1, y2, mouse_rectangle

	if event == cv2.EVENT_LBUTTONDOWN:
		mouse_rectangle = True
		x1,y1 = x,y
	elif event == cv2.EVENT_MOUSEMOVE:
		if mouse_rectangle == True:
			x2,y2 = x, y
	elif event == cv2.EVENT_LBUTTONUP:
		x2,y2 = x, y
		if x2 < x1: x1,x2 = x2,x1
		if y2 < y1: y1,y2 = y2,y1
		mouse_rectangle = False

cv2.namedWindow("Spectrum Analyser", cv2.WINDOW_NORMAL)
cv2.setMouseCallback("Spectrum Analyser", mouse_event)

vc = cv2.VideoCapture(0)

SHEET_DIR = "data/sheet"
sheet_name = "data"

if len([name for name in os.listdir(SHEET_DIR) if os.path.isfile(os.path.join(SHEET_DIR, name))])-1 > 0:
	sheet_name = sheet_name + str(len([name for name in os.listdir(SHEET_DIR) if os.path.isfile(os.path.join(SHEET_DIR, name))])-1)

workbook = xlsxwriter.Workbook('data/sheet/' + sheet_name + '.xlsx')
worksheet = workbook.add_worksheet()
worksheet.set_column("A1:J5", 52)
worksheet.set_default_row(220)
worksheet.write(0,1, "Graph")
worksheet.write(0,2, "Status")
worksheet.write(0,3, "Image")
worksheet.write(0,4, "Settings")
worksheet.write(0,5, "Date")

if row < 0 :
	print("row can't be less then 0!")
	exit()

def AddToWorksheet(row,status):
	worksheet.insert_image(row,1, 'data/graph/' + str(status) + '/' + (str(not_drunkdatanum) if status=="not_drunk" else str(drunkdatanum)) + '.png', {'x_scale': 0.47,'y_scale': 0.5,'x_offset': 2,'y_offset': 2,'positioning': 1})
	worksheet.insert_image(row,3, 'data/frames/' + str(status) + '/' + (str(not_drunkdatanum) if status=="not_drunk" else str(drunkdatanum)) + '.png', {'x_scale': 0.47,'y_scale': 0.5,'x_offset': 2,'y_offset': 2,'positioning': 1})
	worksheet.write(row, 0, row)
	worksheet.write(row, 2, status)
	worksheet.write(row, 4, str(settingsJson))
	worksheet.write(row, 5, str(datetime.datetime.now()))

	print("\n----------------\nSaved!" + "\nRow:" + str(row) + "\nNot Drunk Data Num:" + str(not_drunkdatanum) + "\nDrunk Data Num:" + str(drunkdatanum))


def AutoSpectrumFinder(frame):

	frame = (rotate_bound(frame, angle))
	height, width = frame.shape[:2]

	start_row, start_col = 0, 0
	end_row, end_col = height, width // 2

	frame = frame[start_row:end_row , start_col:end_col]

	grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
	ret,th = cv2.threshold(grayscale,127,255, 0)
	
	contours, hierarchy = cv2.findContours(th, 2, 1)
	cnt = contours
	big_contour = []
	max = 0
	for i in cnt:
		area = cv2.contourArea(i)
		if(area > max):
			max = area
			big_contour = i

	final = cv2.drawContours(frame, big_contour, -1, (0,255,0), 3)
	point = big_contour[int(len(big_contour)/2)][0]

	x = point[0]
	y = point[1]

	return x - 50, y - 20, x + 50, y + 10

if vc.isOpened(): rval, frame = vc.read()
else: rval = False

while rval:
	frame = (rotate_bound(frame, angle))
	prntframe = frame.copy()
	cv2.rectangle(prntframe, (x1, y1), (x2, y2), (0,100,0), 2)
	prntframe = prntframe
	
	try: cutframe = cv2.cvtColor(frame[y1:y2, x1:x2], cv2.COLOR_BGR2GRAY)
	except cv2.error: cutframe = [[0]]
	
	cv2.imshow("Spectrum Analyser", prntframe)
	key = cv2.waitKey(500)
	rval, frame = vc.read()

	graph = [500]

	if plot_graph:
		if not mouse_rectangle:
			if live == False: plot_graph = False
			else: plt.clf()

			for line in list(rotate_clockwise(cutframe)):
				bright = 0
				for point in line: bright += point
				graph.append(bright)

			graph.append(500)
			plt.ylim(top=max(graph))
			plt.ylim(bottom=min(graph))
			if not cmp_graph == [500]:
				plt.plot(cmp_graph[::-1] if flip else cmp_graph, color=(0, 0, 0))
				plt.ylim(top=max([max(cmp_graph), max(graph)]))
				plt.ylim(bottom=min([min(cmp_graph), min(graph)]))
				nor_color = (0, 1, 0)
				for i, x in enumerate(graph): 
					if abs(cmp_graph[i] - x) < sensetivity: nor_color = (1, 0, 0)

			plt.plot(graph[::-1] if flip else graph, color=nor_color)
			plt.draw()
			plt.pause(0.001)
		else: plt.draw()

	if key == ord("w"): y1 -= 2
	elif key == ord("s"): y2 += 2
	elif key == ord("l"): 
		if live: live = False 
		else: live = True
		print("Live = " + str(live))
	elif key == ord("/"):
		print("Searching...")
		try:
			x1, y1, x2, y2 = AutoSpectrumFinder(frame)
			print("Found", (x1, y1, x2, y2))
		except:
			print("Could'nt find, please select manually")
	elif key == ord("\\"): cv2.destroyWindow("mask")
	elif key == ord("a"): x1 -= 2
	elif key == ord("d"): x2 += 2
	elif key == ord("z"): y1 += 2
	elif key == ord("x"): y2 -= 2
	elif key == ord("q"): x1 += 2
	elif key == ord("e"): x2 -= 2
	elif key == ord("r"):
		angle -= 0.5
		if angle <= 0: angle = 360
		print("Angle = " + str(angle))
	elif key == ord("t"):
		angle += 0.5
		if angle >= 360: angle = 0
		print("Angle = " + str(angle))
	elif key == ord("y"):
		angle += 180
		sleep(0.1)
		if angle <= 0: angle = 360
		if angle >= 360: angle = 0
		print("Angle = " + str(angle))
	elif key == ord("n"):
		sensetivity -= 10
		print("sensitivity = " + str(sensetivity))
	elif key == ord("m"):
		sensetivity += 10
		print("sensitivity = " + str(sensetivity))
	elif key == ord("o"):
		plot_graph = False
		plt.clf()
		plt.close()
		print("\nGraph: Closed")
	elif key == ord("p"):
		plot_graph = True
		plt.clf()
		print("\nGraph: Plotting...")
	elif key == ord("v"):
		flip = not flip
		sleep(0.1)
		print("Flip = " + str(flip))
	elif key == ord("c"):
		print("System: Closed")
		save = ""
		while save != 'y' or save != 'n':
			save = input("Do you want to save the sheet? (y/n)")
			if save == "n":
				workbook.close()
				os.remove("data/sheet/" + sheet_name + ".xlsx")
				print("Deleted!")
				break
			else:
				workbook.close()
				break
		break
	elif key == ord("k"):
		cmp_graph = graph
		nor_color = (0,1,0)
	elif key == ord("j"):
		cmp_graph = [500]
		nor_color = (1, 0, 0)
	elif key == ord("h"):
		pickle.dump({"sensitivity": sensetivity, "x1": x1, "y1": y1, "x2": x2, "y2": y2, "angle": angle, "flip": flip}, open("settings.p", "wb"))
	elif key == ord("g"):
		try:
			print("Paused graph, press CTRL+C while focused on the terminal to resume it.")
			plt.pause(10000000)
		except KeyboardInterrupt:
			print("Resumed program")
	elif key == ord("]") and not live:
		print("Recorded Drunk!")
		plt.savefig('data/graph/drunk/' + str(drunkdatanum) + '.png')
		cv2.imwrite("data/frames/drunk/" + str(drunkdatanum) + '.png', cutframe)
		AddToWorksheet(row,"drunk")
		drunkdatanum += 1
		row += 1
	elif key == ord("[") and not live:
		print("Recorded Not Drunk!")
		plt.savefig('data/graph/not_drunk/' + str(not_drunkdatanum) + '.png')
		cv2.imwrite("data/frames/not_drunk/" + str(drunkdatanum) + '.png', cutframe)
		AddToWorksheet(row,"not_drunk")
		not_drunkdatanum += 1
		row += 1
cv2.destroyWindow("Infrared spectrometer")
workbook.close()