OpenCV

Semantic Segmentation_video

728x90
import numpy as np
import pandas as pd
import imutils
import time
import cv2
import os
import matplotlib.pyplot as plt

DEFAULT_FRAME = 1

SET_WIDTH = int(600)
sv = cv2.VideoCapture('data4/video/video.mp4')


try :
  prop = cv2.cv.CV_CAP_PROP_FRAME_COUNT if imutils.is_cv2() else cv2.CAP_PROP_FRAME_COUNT
  total = sv.get(prop)
#   print("[INFO] {} total frames in video.".format(total))
except :
#   print("[INFO] could not determine number of frames in video")
  total = -1

while True :
  grabbed, frame = sv.read()

  if grabbed == False :
    break

  normalize_image = 1 / 255.0
  resize_image_shape = (1024, 512)
  video_frame = imutils.resize(frame, width=SET_WIDTH)
  blob_img = cv2.dnn.blobFromImage(frame, normalize_image, resize_image_shape, 0, 
                                   swapRB = True, crop = False)
  cv_enet_model = cv2.dnn.readNet('data4/enet-cityscapes/enet-model.net')
  cv_enet_model.setInput(blob_img)
  CV_ENET_SHAPE_IMG_COLORS = open('data4/enet-cityscapes/enet-colors.txt').read().split('\n')
  
  # 클래스별로 색정보가 들어있는데, 문자열로 들어있다. 
  # print(CV_ENET_SHAPE_IMG_COLORS)
  
  # 맨 마지막 따옴표를 없애기
  CV_ENET_SHAPE_IMG_COLORS = CV_ENET_SHAPE_IMG_COLORS[ : -2+1]
  
  
  temp = []
  for color in CV_ENET_SHAPE_IMG_COLORS:
    color_list = color.split(',')      
    color_num_list = np.array(color_list).astype('int')  
    print(color_num_list)
    temp.append(color_num_list)
  
  CV_ENET_SHAPE_IMG_COLORS = np.array(temp)

  # 모델이, 세그멘테이션 추론(예측)하는데 얼마나 걸렸는지 측정.
  start_time = time.time()
  cv_enet_model_output = cv_enet_model.forward()
  end_time = time.time()

  (classes_num, height, width) = cv_enet_model_output.shape[1:4]

  class_map = np.argmax(cv_enet_model_output[0], axis=0)

  mask_class_map = CV_ENET_SHAPE_IMG_COLORS[class_map]

  mask_class_map = cv2.resize(mask_class_map, (video_frame.shape[1], video_frame.shape[0]) ,
             interpolation = cv2.INTER_NEAREST)
  
  cv_enet_model_output = ( (0.3 * video_frame) + (0.7 * mask_class_map) ).astype('uint8')
  
  cv2.imshow('frame', cv_enet_model_output)
  if cv2.waitKey(25)&0xFF==27:
      break



sv.release()
cv2.destroyAllWindow()
        
728x90

'OpenCV' 카테고리의 다른 글

road_marking_detection  (0) 2021.04.20
find contures  (0) 2021.04.20
Canny  (0) 2021.04.20
sharpen  (0) 2021.04.20
laplacian  (0) 2021.04.20