project.py 3.67 KB
"""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-image.shape[1]/2)+ abs(y-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);