project.py
3.69 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
"""Demo for use yolo v3
"""
import os
import time
import cv2
import numpy as np
from model.yolo_model import YOLO
def process_image(img):
"""Resize, reduce and expand image.
# Argument:
img: original image.
# Returns
image: ndarray(64, 64, 3), processed image.
"""
image = cv2.resize(img, (416, 416),
interpolation=cv2.INTER_CUBIC)
image = np.array(image, dtype='float32')
image /= 255.
image = np.expand_dims(image, axis=0)
return image
def get_classes(file):
"""Get classes name.
# Argument:
file: classes name for database.
# Returns
class_names: List, classes name.
"""
with open(file) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
if __name__ == '__main__':
# 파일 열기
camera = cv2.VideoCapture("input/example.mp4")
# Yolo 학습
yolo = YOLO(0.6, 0.5)
file = 'data/coco_classes.txt'
all_classes = get_classes(file)
# 1 카운트 할 때마다 frame 얻어서 파일로 저장
success,image = camera.read()
count = 0
while success:
cv2.imwrite("mid/frame%d.png" % count, image) # save frame as JPEG file
success,image = camera.read()
count += 1
# 각 프레임 별로 Image Detection 후 프레임 번호, 객체 이름(name)과 객체의 크기(size), 객체가 얼마나 가운데 있는지(coordinatevalue) 저장
detectionInfo = []
for i in range(count):
#filename = "mid/frame"+str(i)+"."
image = cv2.imread("mid/frame%d.png" % i)
pimage = process_image(image)
boxes, classes, scores = yolo.predict(pimage, image.shape)
for box, score, cl in zip(boxes, scores, classes):
x, y, w, h = box
name = all_classes[cl]
size = w*h
# 얼마나 가운데인지 확인하는 알고리즘
coordinatevalue = abs((x+0.5*w)-image.shape[1]/2) + abs((y-h*0.5)-image.shape[0]/2)/(image.shape[0]/image.shape[1])
detectionInfo.append([i, name, size, coordinatevalue])
# name 별로 크기가 가장 크거나 물체가 프레임의 가운데 있는 프레임 번호 목록 얻어오기
indexlist = []
# 검출된 물체 리스트(중복 없이)
namelist = set()
for i in range(count):
namelist.add(detectionInfo[i][1])
namelist = list(namelist)
# 검출된 물체 중 가장 크기가 큰 것을 불러옴
for i in range(len(namelist)):
maxindex = 0
maxvalue = 0
for j in range(count):
if detectionInfo[j][1] == namelist[i]: # 물체 리스트 중에서
if detectionInfo[j][2] > maxvalue: # 가장 큰 값을 갱신한다면
maxvalue = detectionInfo[j][2]
maxindex = j
indexlist.append(maxindex)
# 검출된 물체 중 가장 가운데 있는 것을 불러
for i in range(len(namelist)):
minindex = 0
minvalue = 2000
for j in range(count):
if detectionInfo[j][1] == namelist[i]: # 물체 리스트 중에서
if detectionInfo[j][3] < minvalue: # 가장 큰 값을 갱신한다면
minvalue = detectionInfo[j][3]
minindex = j
indexlist.append(minindex)
#얻어온 프레임 목록 결과값으로 저장, 선명도 높이는 작업 수행
for i in range(len(indexlist)):
image = cv2.imread("mid/frame%d.png" % indexlist[i])
output1 = cv2.GaussianBlur(image,(5,5),0)
#kernel_sharpen_1 = np.array([[-1,-1,-1],[-1,9,-1],[-1,-1,-1]])
#output2 = cv2.filter2D(output1,-1,kernel_sharpen_1)
cv2.imwrite("output/output%d.png"% (i+1), output1);