728x90
< 전체 코드 >
더보기
더보기
import tensorflow as tf
import os
import pathlib
import numpy as np
import zipfile
import matplotlib.pyplot as plt
from PIL import Image
import cv2
from object_detection.utils import ops as utils_ops
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as viz_utils
import time
# 내 로컬에 설치된 레이블 파일을, 인덱스와 연결시킨다.
PATH_TO_LABELS = 'C:\\Users\\5-20\\Documents\\TensorFlow\\models\\research\\object_detection\\data\\mscoco_label_map.pbtxt'
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS)
# 모델 로드하는 함수.
# https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md
# 위의 사이트에서 모델을 가져올수있다.
# /20200711/efficientdet_d0_coco17_tpu-32.tar.gz
# /20200711/ssd_mobilenet_v2_320x320_coco17_tpu-8.tar.gz
# /20210210/centernet_mobilenetv2fpn_512x512_coco17_od.tar.gz
# /20200711/efficientdet_d0_coco17_tpu-32.tar.gz
# Download and extract model
def download_model(model_name, model_date):
base_url = 'http://download.tensorflow.org/models/object_detection/tf2/'
model_file = model_name + '.tar.gz'
model_dir = tf.keras.utils.get_file(fname=model_name,
origin=base_url + model_date + '/' + model_file,
untar=True)
return str(model_dir)
MODEL_DATE = '20200711'
MODEL_NAME = 'ssd_mobilenet_v2_320x320_coco17_tpu-8'
PATH_TO_MODEL_DIR = download_model(MODEL_NAME, MODEL_DATE)
def load_model(model_dir) :
model_full_dir = model_dir + "/saved_model"
# Load saved model and build the detection function
detection_model = tf.saved_model.load(model_full_dir)
return detection_model
detection_model = load_model(PATH_TO_MODEL_DIR)
def show_inference(detection_model, image_np) :
# The input needs to be a tensor, convert it using `tf.convert_to_tensor`.
input_tensor = tf.convert_to_tensor(image_np)
# The model expects a batch of images, so add an axis with `tf.newaxis`.
input_tensor = input_tensor[tf.newaxis, ...]
# input_tensor = np.expand_dims(image_np, 0)
detections = detection_model(input_tensor)
# print(detections)
# All outputs are batches tensors.
# Convert to numpy arrays, and take index [0] to remove the batch dimension.
# We're only interested in the first num_detections.
num_detections = int(detections.pop('num_detections'))
detections = {key: value[0, :num_detections].numpy()
for key, value in detections.items()}
detections['num_detections'] = num_detections
# detection_classes should be ints.
detections['detection_classes'] = detections['detection_classes'].astype(np.int64)
print(detections)
image_np_with_detections = image_np.copy()
viz_utils.visualize_boxes_and_labels_on_image_array(
image_np_with_detections,
detections['detection_boxes'],
detections['detection_classes'],
detections['detection_scores'],
category_index,
use_normalized_coordinates=True,
max_boxes_to_draw=200,
min_score_thresh=.30,
agnostic_mode=False)
cv2.imshow( 'result' , image_np_with_detections)
def save_inference(detection_model, image_np, video_writer):
# The input needs to be a tensor, convert it using `tf.convert_to_tensor`.
input_tensor = tf.convert_to_tensor(image_np)
# The model expects a batch of images, so add an axis with `tf.newaxis`.
input_tensor = input_tensor[tf.newaxis, ...]
# input_tensor = np.expand_dims(image_np, 0)
detections = detection_model(input_tensor)
# print(detections)
# All outputs are batches tensors.
# Convert to numpy arrays, and take index [0] to remove the batch dimension.
# We're only interested in the first num_detections.
num_detections = int(detections.pop('num_detections'))
detections = {key: value[0, :num_detections].numpy()
for key, value in detections.items()}
detections['num_detections'] = num_detections
# detection_classes should be ints.
detections['detection_classes'] = detections['detection_classes'].astype(np.int64)
print(detections)
image_np_with_detections = image_np.copy()
viz_utils.visualize_boxes_and_labels_on_image_array(
image_np_with_detections,
detections['detection_boxes'],
detections['detection_classes'],
detections['detection_scores'],
category_index,
use_normalized_coordinates=True,
max_boxes_to_draw=200,
min_score_thresh=.30, # 예측한 네모박스가 적어도 확신이 30% 이상 되는것만 박스를 그린다
agnostic_mode=False)
video_writer.write(image_np_with_detections)
# 비디오를 실행하는 코드
cap = cv2.VideoCapture('data/dashcam2.mp4')
# cap = cv2.VideoCapture(0) # 카메라 캠
if cap.isOpened() == False :
print('비디오 실행에러')
else :
frame_width = int( cap.get(3) )
frame_height = int( cap.get(4) )
out = cv2.VideoWriter('data/output.avi', cv2.VideoWriter_fourcc('M','J','P','g'), 20, (frame_width, frame_height))
# 비디오 캡쳐해서 이미지를 1장씩 가져온다
# 이 1장의 이미지를 오브젝트 디텍션 한다
while cap.isOpened() :
ret, frame = cap.read()
if ret == True :
# frame 이 이미지에 대한 넘파이 어레이 이므로
# 이 frame 을 오브젝트 디텔션 한다.
# 학습용으로, 동영상으로 저장하는 코드를
# 수정하세요
start_time = time.time()
# 동영상을 디텍션 한 후, 파일로 저장하는 함수
save_inference(detection_model, frame, out)
# 돋영상을 실시간으로 바로 화면에서 디텍팅하는것
# show_inference(detection_model, frame)
end_time = time.time()
print('연산에 걸린시간', str(end_time-start_time))
if cv2.waitKey(27) & 0xFF == 27 :
break
else :
break
cap.release()
out.release()
cv2.destroyAllWindows()
비디오를 실행하는 코드부분에서 동영상 저장코드를 추가한다.
# 비디오를 실행하는 코드부분
cap = cv2.VideoCapture('data/dashcam2.mp4')
# cap = cv2.VideoCapture(0) # 카메라 캠
if cap.isOpened() == False :
print('비디오 실행에러')
else :
frame_width = int( cap.get(3) )
frame_height = int( cap.get(4) )
out = cv2.VideoWriter('data/output.avi', cv2.VideoWriter_fourcc('M','J','P','g'), 20, (frame_width, frame_height))
# 비디오 캡쳐해서 이미지를 1장씩 가져온다
# 이 1장의 이미지를 오브젝트 디텍션 한다
while cap.isOpened() :
ret, frame = cap.read()
if ret == True :
# frame 이 이미지에 대한 넘파이 어레이 이므로
# 이 frame 을 오브젝트 디텔션 한다.
# 학습용으로, 동영상으로 저장하는 코드를
# 수정하세요
start_time = time.time()
# 동영상을 디텍션 한 후, 파일로 저장하는 함수
save_inference(detection_model, frame, out)
# 돋영상을 실시간으로 바로 화면에서 디텍팅하는것
# show_inference(detection_model, frame)
end_time = time.time()
print('연산에 걸린시간', str(end_time-start_time))
if cv2.waitKey(27) & 0xFF == 27 :
break
else :
break
cap.release()
out.release()
cv2.destroyAllWindows()< 추가한 코드 부분 >