project.py
3.38 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
import time
import cv2
import numpy as np
from model.yolo_model import YOLO
def process_image(img):
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):
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/%d.%s.png"% (i+1,detectionInfo[indexlist[i]][1]), output1);