SSD lufficc OpenCV WebCam Demo

camera_demo.py

import glob
import os
import time

import torch
from PIL import Image
from vizer.draw import draw_boxes

import cv2

from ssd.config import cfg
from ssd.data.datasets import COCODataset, VOCDataset
import argparse
import numpy as np

from ssd.data.transforms import build_transforms
from ssd.modeling.detector import build_detection_model
from ssd.utils import mkdir
from ssd.utils.checkpoint import CheckPointer

#from torch2trt import torch2trt

#COCO Dataset Type class_names
#TODO REPLACE
#class_names = COCODataset.class_names
class_names = ['__background__','parking','stop','mertic']

#@torch.no_grad()

#cfg.merge_from_file("configs/vgg_ssd300_coco_trainval35k.yaml")
cfg.merge_from_file("configs/mobilenet_v2_ssd320_coco712.yaml")
cfg.OUTPUT_DIR = "outputs/mobilenet_v2_ssd320_coco0712"
weights_path = "outputs/mobilenet_v2_ssd320_coco0712/model_final.pth"
cfg.MODEL.NUM_CLASSES = 3
#cfg.merge_from_list(args.opts)
cfg.freeze()
print(cfg)
ckpt=None
score_threshold=0.5
dataset_type="coco"


device = torch.device(cfg.MODEL.DEVICE)
model = build_detection_model(cfg)
model = model.to(device)
print(cfg.OUTPUT_DIR)

checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
checkpointer.load(ckpt, use_latest=ckpt is None)
weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
print('Loaded weights from {}'.format(weight_file))

cpu_device = torch.device("cpu")
transforms = build_transforms(cfg, is_train=False)
model.eval()

#data = torch.randn((1, 3, 300, 300)).cuda()
#model_trt = torch2trt(model, [data])#, fp16_mode=True)

cap = cv2.VideoCapture(0)
#cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
#cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

while True:
    __, frame = cap.read()
    #frame = cv2.GaussianBlur(frame,(15,15),0.9)
    height, width = frame.shape[:2]
    images = transforms(frame)[0].unsqueeze(0)
    #print(images)
    
    with torch.no_grad():
        start = time.time()
        result = model(images.to(device))[0]
        inference_time = time.time() - start
        print("Timp procesare inferenta {0:.2f} ms".format(inference_time*100))
        result = result.resize((width, height)).to(cpu_device).numpy()
        boxes, labels, scores = result['boxes'], result['labels'], result['scores']

    indices = scores > score_threshold
    boxes = boxes[indices]
    labels = labels[indices]
    scores = scores[indices]

    drawn_image = draw_boxes(frame, boxes, labels, scores, class_names).astype(np.uint8)
    cv2.imshow("Detection",drawn_image)
    #cv2.imshow("camera", frame)
    
    if(cv2.waitKey(1) == ord("q")):
        cap.release()
        break