jetson ssd

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jun 24 15:38:48 2021

@author: ledi
"""

import cv2


def gstreamer_pipeline(
    capture_width=1280,
    capture_height=720,
    display_width=1280,
    display_height=720,
    framerate=60,
    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,
        )
    )



"""trt_ssd.py

This script demonstrates how to do real-time object detection with
TensorRT optimized Single-Shot Multibox Detector (SSD) engine.
"""


import time
import argparse

# import cv2
import pycuda.autoinit  # This is needed for initializing CUDA driver

from utils.ssd_classes import get_cls_dict
from utils.ssd import TrtSSD
from utils.camera import add_camera_args, Camera
from utils.display import open_window, set_display, show_fps
from utils.visualization import BBoxVisualization


WINDOW_NAME = 'TrtSsdDemo'
INPUT_HW = (300, 300)
SUPPORTED_MODELS = [
    'ssd_mobilenet_v1_coco',
    'ssd_mobilenet_v1_egohands',
    'ssd_mobilenet_v2_coco',
    'ssd_mobilenet_v2_egohands',
    'ssd_inception_v2_coco',
    'ssdlite_mobilenet_v2_coco',
]


def parse_args():
    """Parse input arguments."""
    desc = ('Capture and display live camera video, while doing '
            'real-time object detection with TensorRT optimized '
            'SSD model on Jetson Nano')
    parser = argparse.ArgumentParser(description=desc)
    parser = add_camera_args(parser)
    parser.add_argument('-m', '--model', type=str,
                        default='ssd_mobilenet_v1_coco',
                        choices=SUPPORTED_MODELS)
    args = parser.parse_args()
    return args









model_path='model ssd_mobilenet_v1_coco'
cls_dict = get_cls_dict(model_path.split('_')[-1])
trt_ssd = TrtSSD('ssd_mobilenet_v1_coco', (300,300))

open_window(
    WINDOW_NAME, 'Camera TensorRT SSD Demo',
    1080, 1440)
vis = BBoxVisualization(cls_dict)
# loop_and_detect(cam, trt_ssd, conf_th=0.3, vis=vis)

cap = cv2.VideoCapture(gstreamer_pipeline(flip_method=0), cv2.CAP_GSTREAMER)


full_scrn = False
fps = 0.0
tic = time.time()

while cap.isOpened():
    flag, img = cap.read()
    
    boxes, confs, clss = trt_ssd.detect(img, 0.3)
    img = vis.draw_bboxes(img, boxes, confs, clss)
    img = show_fps(img, fps)
    cv2.imshow(WINDOW_NAME, img)
    toc = time.time()
    curr_fps = 1.0 / (toc - tic)
    # calculate an exponentially decaying average of fps number
    fps = curr_fps if fps == 0.0 else (fps*0.95 + curr_fps*0.05)
    tic = toc
    key = cv2.waitKey(1)
    if key == 27:  # ESC key: quit program
        break
    elif key == ord('F') or key == ord('f'):  # Toggle fullscreen
        full_scrn = not full_scrn
        set_display(WINDOW_NAME, full_scrn)




# cap.release()
# cv2.destroyAllWindows()


# if __name__ == '__main__':
#     main()

# if __name__ == "__main__":
#     show_camera()