project.py 3.94 KB
import time
import math
import cv2
import numpy as np
from model.yolo_model import YOLO

class Point2D:
    def __init__(self, width, height):
        self.width = width
        self.height = height

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
		
            # 얼마나 가운데인지 확인하는 알고리즘
            object = Point2D(width= x + h/2, height= y + w/2)
            a = image.shape[1]/2 - object.width
            b = image.shape[0]/2 - object.height
            coordinatevalue = math.sqrt((a*a)+(b*b))
            # 객체 정보 및 계산 값 저장
            detectionInfo.append([i, name, size, coordinatevalue])

    for i in range(len(detectionInfo)):
        print(detectionInfo[i])

    # 검출된 물체 리스트(중복 없이)
    namelist = {}
    for i in range(len(detectionInfo)):
        if not detectionInfo[i][1] in namelist:
            namelist[detectionInfo[i][1]] = []

    # 계획1 : 뒤 두 알고리즘 일정 비율로 합치기
    # 크기
    for objectName in namelist.keys():
        maxindex = 0
        maxvalue = 0
        for j in range(len(detectionInfo)):
            if detectionInfo[j][1] == objectName:
                if detectionInfo[j][2] > maxvalue:
                    maxvalue = detectionInfo[j][2]
                    maxindex = detectionInfo[j][0]
        namelist[objectName].append(maxindex)
   
    for objectname, framelist in namelist.items():
        image = cv2.imread("mid/frame%d.png" % framelist[0])
        output1 = cv2.GaussianBlur(image, (5,5), 0)
        cv2.imwrite("output1/%s.png"% (objectname), output1)    

	# 가운데 위치
    for objectName in namelist.keys():
        namelist[objectName] = []

    for objectName in namelist.keys():
        minindex = 0
        minvalue = 999999
        for j in range(len(detectionInfo)):
            if detectionInfo[j][1] == objectName:
                if detectionInfo[j][3] < minvalue:
                    minvalue = detectionInfo[j][3]
                    minindex = detectionInfo[j][0]
        namelist[objectName].append(minindex)
    
    for objectname, framelist in namelist.items():
        image = cv2.imread("mid/frame%d.png" % framelist[0])
        output2 = cv2.GaussianBlur(image, (5,5), 0)
        cv2.imwrite("output2/%s.png"% (objectname), output2)


# 계획2 : 프레임별로 나온 객체 겹치는 부분 제외하고 넓이 구해 큰거 Indexlist에 넣기

# 계획3 : 객체가 특정 위치에 있는 프레임 뽑기








#kernel_sharpen_1 = np.array([[-1,-1,-1],[-1,9,-1],[-1,-1,-1]])
#output2 = cv2.filter2D(output1,-1,kernel_sharpen_1)