import jetson.inference
import jetson.utils
import cv2
import numpy as np

face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

net = jetson.inference.detectNet(argv=['--model=ssd-mobilenet.onnx', '--labels=labels.txt', '--input-blob=input_0', '--output-cvg=scores', '--output-bbox=boxes', '--overlay=box,labels,conf'], threshold=80)
#camera = jetson.utils.videoSource("/dev/video0")      # '/dev/video0' for V4L2
display = jetson.utils.videoOutput("display://0") # 'my_video.mp4' for file

cap = cv2.VideoCapture(0)

while display.IsStreaming():

	#img = camera.Capture()
	_, img = cap.read()
	image = cv2.cvtColor(img, cv2.COLOR_BGR2RGBA).astype(np.float32)
	image = jetson.utils.cudaFromNumpy(image)
	detections = net.Detect(image, 640, 480)
	display.Render(image)
	display.SetStatus("Object Detection | Network {:.0f} FPS".format(net.GetNetworkFPS()))

	#detections = net.Detect(img)
	image = jetson.utils.cudaToNumpy(image, 640, 480, 4)
	image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB).astype(np.uint8)
	
	# Convert into grayscale
	gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
	# Detect faces
	faces = face_cascade.detectMultiScale(gray, 1.1, 4)

	# Draw rectangle around the faces
	for (x, y, w, h) in faces:
	    cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)

	# Display the output
	cv2.imshow('img', img)
	#cv2.imshow('image', image)



	k = cv2.waitKey(30) & 0xff
	if k==27:
		break
        
# Release the VideoCapture object
cap.release()
cv2.destroyAllWindows()