# This is an bottle detection system for use by the Clorox sponsored project at KSU,
# nicknamed Blipper (Bottle Flipper). Output GPIO signals are sent from the 
# Jetson Nano to the PLC controlling the robotic arm. 
# Object detection code by NVIDIA, edited and expanded for this project by
# Preston Delaware. April 15, 2021

import jetson.inference
import jetson.utils
import Jetson.GPIO as GPIO

# define constants for object detection
L24 = 1
L43 = 2
R24 = 3
R43 = 4

# GPIO pin definitions
detectEnable = 8      # Pin 24 BCM pin 8 [white wire]
bottleType = 18       # Pin 12, BCM pin 18 [blue wire]
directionOutput = 20  # Pin 38, BCM pin 20 [purple wire]

# Pin Setup:
GPIO.setmode(GPIO.BCM)
# set pin as an output pin with optional initial state of LOW
GPIO.setup(detectEnable, GPIO.OUT, initial=GPIO.LOW)
GPIO.setup(directionOutput, GPIO.OUT, initial=GPIO.LOW)
GPIO.setup(bottleType, GPIO.OUT, initial=GPIO.LOW)

############################
# output code is as follows
# if detectEnable == 1:
#   0  -  Left-facing
#   1  -  Right-facing
# if detectEnable == 0:
#   do nothing
#
# bottleType designation
#   1  =  24 oz
#   0  =  43 oz
############################

# load the object detection model
net = jetson.inference.detectNet(argv=['--model=jetson-inference/python/training/detection/ssd/models/bottles/ssd-mobilenet.onnx', '--labels=jetson-inference/python/training/detection/ssd/models/bottles/labels.txt', '--input-blob=input_0', '--output-cvg=scores', '--output-bbox=boxes', '--threshold=0.8'])

# set up camera
camera = jetson.utils.videoSource("csi://0")      # '/dev/video0'

# set up display for output to screen
display = jetson.utils.videoOutput("display://0")

detectL24Count   = 0			# create variables for stability
detectL43Count   = 0
detectR24Count   = 0
detectR43Count   = 0
noDetectionCount = 0
detectionCount   = 0
maxCount 	 = 5
lastDetection    = 0			# empty variable for now

while display.IsStreaming():		# while display window is open
	frame = camera.Capture()		# take incoming video frame
	detection = net.Detect(frame)	# detect objects in the frame and save to detection variable

	display.Render(frame)		# show the frame
	display.SetStatus("Object Detection | Network {:.0f} FPS".format(net.GetNetworkFPS()))	# show FPS of camera
	print(detection)	# for debugging

	# if L24 bottle detected for several frames in a row
	if detectL24Count >= maxCount:
		detectL24Count = 0
             	# send Left facing signal
		GPIO.output(directionOutput, GPIO.HIGH)
		# send 24oz signal
		GPIO.output(bottleType, GPIO.HIGH)

	# if L43 bottle detected for several frames in a row
	if detectL43Count >= maxCount:
		detectL43Count = 0
		# send Left facing signal
		GPIO.output(directionOutput, GPIO.HIGH)
		# send 43oz signal
		GPIO.output(bottleType, GPIO.LOW)

	# if R24 bottle detected for several frames in a row
	if detectR24Count >= maxCount:
		detectR24Count = 0
             	# send Right facing signal
		GPIO.output(directionOutput, GPIO.LOW)
		# send 24oz signal
		GPIO.output(bottleType, GPIO.HIGH)

	# if R43 bottle detected for several frames in a row
	if detectR43Count >= maxCount:
		detectR43Count = 0
             	# send Right facing signal
		GPIO.output(directionOutput, GPIO.LOW)
		# send 43oz signal
		GPIO.output(bottleType, GPIO.LOW)

	# if no bottle detected after several frames...
	if noDetectionCount >= maxCount:
		noDetectionCount = 0			# turn off all signals
		detectL24Count = 0
		detectL43Count = 0
		detectR24Count = 0
		detectR43Count = 0
		GPIO.output(detectEnable, GPIO.LOW) 
		GPIO.output(directionOutput, GPIO.LOW)
		GPIO.output(bottleType, GPIO.LOW)

	# if we have detected an object for several frames...
	if detectionCount >= maxCount:
		detectionCount = 0		
		GPIO.output(detectEnable, GPIO.HIGH)    # detectEnable = 1, tell robot to pay attention

	# if detection empty (nothing detected)
	if detection == []:		
		noDetectionCount += 1		# increment count		

	# if an object was detected
	if detection != []:		
		detectionCount += 1		# increment count
		# if inconsistent detections
		if detection[0].ClassID != lastDetection:
			detectL24Count = 0			# reset all counts to zero
			detectL43Count = 0
			detectR24Count = 0
			detectR43Count = 0	
		if detection[0].ClassID == L24: # if ClassID = left 24oz
			detectL24Count += 1	# increment count
		if detection[0].ClassID == L43: # if ClassID = left 43oz
			detectL43Count += 1     # increment count      		
		if detection[0].ClassID == R24: # if ClassID = right 24oz
			detectR24Count += 1	# increment count
		if detection[0].ClassID == R43: # if ClassID = right 43oz
			detectR43Count += 1	# increment count
		lastDetection = detection[0].ClassID	# save the current detection

# end while loop
GPIO.output(detectEnable, GPIO.LOW) # turn off all pins when exiting program
GPIO.output(directionOutput, GPIO.LOW)
GPIO.output(bottleType, GPIO.LOW)