# MIT License
# Copyright (c) 2019 JetsonHacks
# See license
# Using a CSI camera (such as the Raspberry Pi Version 2) connected to a 
# NVIDIA Jetson Nano Developer Kit using OpenCV
# Drivers for the camera and OpenCV are included in the base image
from __future__ import print_function

import os
import argparse
import numpy as np
import cv2

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

from train import Net

# gstreamer_pipeline returns a GStreamer pipeline for capturing from the CSI camera
# Defaults to 1280x720 @ 60fps 
# Flip the image by setting the flip_method (most common values: 0 and 2)
# display_width and display_height determine the size of the window on the screen
def gstreamer_pipeline (capture_width=1280, capture_height=720, display_width=640, display_height=360, framerate=20, flip_method=0) :   
    return ('nvarguscamerasrc ! ' 
    'video/x-raw(memory:NVMM), '
    'width=(int)%d, height=(int)%d, '
    'format=(string)NV12, framerate=(fraction)%d/1 ! '
    'nvvidconv flip-method=%d ! '
    'video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! '
    'videoconvert ! '
    'video/x-raw, format=(string)BGR ! appsink'  % (capture_width,capture_height,framerate,flip_method,display_width,display_height))


def transform_inputs(image, device):
    img = cv2.resize(image, (28, 28), interpolation=cv2.INTER_AREA)
    # img = cv2.resize(image, (28, 28))
    # image = cv2.threshold(image0, 50, 255, cv2.THRESH_BINARY)[1]
    blur = cv2.GaussianBlur(img,(5,5),0)
    image = cv2.threshold(blur,0,255,cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)[1]

    # image = image[np.newaxis, ...]
    transform = transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                ])
    inputs = transform(image)
    inputs = inputs.unsqueeze(0)
    return inputs.to(device)


def show_camera(args):
    # To flip the image, modify the flip_method parameter (0 and 2 are the most common)
    print(gstreamer_pipeline(flip_method=0))
    labels = []
    with open("data/labels.txt", "r") as f:
        for line in f:
            labels.append(line.strip())

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    model = Net(n_classes=len(labels))
    model.load_state_dict(torch.load(args.model_path))
    model = model.to(device)

    font = cv2.FONT_HERSHEY_SIMPLEX
    fontScale = 1
    org = (30, 50)
    color = (0, 0, 255)
    thickness = 2

    cap = cv2.VideoCapture(gstreamer_pipeline(flip_method=0), cv2.CAP_GSTREAMER)
    if cap.isOpened():
        window_handle = cv2.namedWindow('Camera', cv2.WINDOW_AUTOSIZE)

        # Window 
        while cv2.getWindowProperty('Camera',0) >= 0:
            ret_val, img = cap.read()

            # Convert to grayscale and apply Gaussian filtering
            im_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            im_gray = cv2.GaussianBlur(im_gray, (5, 5), 0)

            # Threshold the image
            ret, im_th = cv2.threshold(im_gray, 90, 255, cv2.THRESH_BINARY_INV)

            # Find contours in the image
            im2, ctrs, hier = cv2.findContours(im_th, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # Get rectangles contains each contour
            rects = [cv2.boundingRect(ctr) for ctr in ctrs]

            # For each rectangular region, calculate HOG features and predict
            # the digit using Linear SVM.
            for rect in rects:
                # Draw the rectangles
                cv2.rectangle(img, (rect[0], rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), (0, 255, 0), 3) 
                # Make the rectangular region around the digit
                leng = int(rect[3] * 1.6)
                pt1 = int(rect[1] + rect[3] // 2 - leng // 2)
                pt2 = int(rect[0] + rect[2] // 2 - leng // 2)
                roi = im_gray[pt1:pt1+leng, pt2:pt2+leng]
                # Resize the image
                h, w = roi.shape
                if h > 10 and w > 10:
                    # Transform inputs
                    inputs = transform_inputs(roi, device)
                    # Run Model Evaluation
                    output = model(inputs)
                    result = output.data.cpu().numpy().argmax()
                    cv2.putText(img, labels[result], (rect[0], rect[1]),cv2.FONT_HERSHEY_DUPLEX, 2, (0, 255, 255), 3)
            
            cv2.imshow("Camera", img)

	    # This also acts as 
            keyCode = cv2.waitKey(30) & 0xff
            # Stop the program on the ESC key
            if keyCode == 27:
               break
        cap.release()
        cv2.destroyAllWindows()
    else:
        print('Unable to open camera')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--model_path', default="models/model.pt")
    args = parser.parse_args()
    show_camera(args)

# MIT License
# Copyright (c) 2019 JetsonHacks
# See license
# Using a CSI camera (such as the Raspberry Pi Version 2) connected to a 
# NVIDIA Jetson Nano Developer Kit using OpenCV
# Drivers for the camera and OpenCV are included in the base image
from __future__ import print_function

import os
import argparse
import numpy as np
import cv2

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

from train import Net

# gstreamer_pipeline returns a GStreamer pipeline for capturing from the CSI camera
# Defaults to 1280x720 @ 60fps 
# Flip the image by setting the flip_method (most common values: 0 and 2)
# display_width and display_height determine the size of the window on the screen
def gstreamer_pipeline (capture_width=1280, capture_height=720, display_width=640, display_height=360, framerate=20, flip_method=0) :   
    return ('nvarguscamerasrc ! ' 
    'video/x-raw(memory:NVMM), '
    'width=(int)%d, height=(int)%d, '
    'format=(string)NV12, framerate=(fraction)%d/1 ! '
    'nvvidconv flip-method=%d ! '
    'video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! '
    'videoconvert ! '
    'video/x-raw, format=(string)BGR ! appsink'  % (capture_width,capture_height,framerate,flip_method,display_width,display_height))


def transform_inputs(image, device):
    img = cv2.resize(image, (28, 28), interpolation=cv2.INTER_AREA)
    # img = cv2.resize(image, (28, 28))
    # image = cv2.threshold(image0, 50, 255, cv2.THRESH_BINARY)[1]
    blur = cv2.GaussianBlur(img,(5,5),0)
    image = cv2.threshold(blur,0,255,cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)[1]

    # image = image[np.newaxis, ...]
    transform = transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                ])
    inputs = transform(image)
    inputs = inputs.unsqueeze(0)
    return inputs.to(device)


def show_camera(args):
    # To flip the image, modify the flip_method parameter (0 and 2 are the most common)
    print(gstreamer_pipeline(flip_method=0))
    labels = []
    with open("data/labels.txt", "r") as f:
        for line in f:
            labels.append(line.strip())

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    model = Net(n_classes=len(labels))
    model.load_state_dict(torch.load(args.model_path))
    model = model.to(device)

    font = cv2.FONT_HERSHEY_SIMPLEX
    fontScale = 1
    org = (30, 50)
    color = (0, 0, 255)
    thickness = 2

    cap = cv2.VideoCapture(gstreamer_pipeline(flip_method=0), cv2.CAP_GSTREAMER)
    if cap.isOpened():
        window_handle = cv2.namedWindow('Camera', cv2.WINDOW_AUTOSIZE)

        # Window 
        while cv2.getWindowProperty('Camera',0) >= 0:
            ret_val, img = cap.read()

            # Convert to grayscale and apply Gaussian filtering
            im_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            im_gray = cv2.GaussianBlur(im_gray, (5, 5), 0)

            # Threshold the image
            ret, im_th = cv2.threshold(im_gray, 90, 255, cv2.THRESH_BINARY_INV)

            # Find contours in the image
            im2, ctrs, hier = cv2.findContours(im_th, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # Get rectangles contains each contour
            rects = [cv2.boundingRect(ctr) for ctr in ctrs]

            # For each rectangular region, calculate HOG features and predict
            # the digit using Linear SVM.
            for rect in rects:
                # Draw the rectangles
                cv2.rectangle(img, (rect[0], rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), (0, 255, 0), 3) 
                # Make the rectangular region around the digit
                leng = int(rect[3] * 1.6)
                pt1 = int(rect[1] + rect[3] // 2 - leng // 2)
                pt2 = int(rect[0] + rect[2] // 2 - leng // 2)
                roi = im_gray[pt1:pt1+leng, pt2:pt2+leng]
                # Resize the image
                h, w = roi.shape
                if h > 10 and w > 10:
                    # Transform inputs
                    inputs = transform_inputs(roi, device)
                    # Run Model Evaluation
                    output = model(inputs)
                    result = output.data.cpu().numpy().argmax()
                    cv2.putText(img, labels[result], (rect[0], rect[1]),cv2.FONT_HERSHEY_DUPLEX, 2, (0, 255, 255), 3)
            
            cv2.imshow("Camera", img)

	    # This also acts as 
            keyCode = cv2.waitKey(30) & 0xff
            # Stop the program on the ESC key
            if keyCode == 27:
               break
        cap.release()
        cv2.destroyAllWindows()
    else:
        print('Unable to open camera')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--model_path', default="models/model.pt")
    args = parser.parse_args()
    show_camera(args)

# MIT License # Copyright (c) 2019 JetsonHacks # See license # Using a CSI camera (such as the Raspberry Pi Version 2) connected to a # NVIDIA Jetson Nano Developer Kit using OpenCV # Drivers for the camera and OpenCV are included in the base image # gstreamer_pipeline returns a GStreamer pipeline for capturing from the CSI camera # Defaults to 1280x720 @ 60fps # Flip the image by setting the flip_method (most common values: 0 and 2) # display_width and display_height determine the size of the window on the screen # img = cv2.resize(image, (28, 28)) # image = cv2.threshold(image0, 50, 255, cv2.THRESH_BINARY)[1] # image = image[np.newaxis, ...] # To flip the image, modify the flip_method parameter (0 and 2 are the most common) # Window # Convert to grayscale and apply Gaussian filtering # Threshold the image # Find contours in the image # Get rectangles contains each contour # For each rectangular region, calculate HOG features and predict # the digit using Linear SVM. # Draw the rectangles # Make the rectangular region around the digit # Resize the image # Transform inputs # Run Model Evaluation # This also acts as # Stop the program on the ESC key