본문 바로가기

Study/class note

딥러닝 / 강아지 품종 detection

import matplotlib.pyplot as plt
import cv2
import numpy as np
from IPython.display import display, Image
import ipywidgets as widgets
import threading
# Stop button
# ================
stopButton = widgets.ToggleButton(
    value=False,
    description='Stop',
    disabled=False,
    button_style='danger', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Description',
    icon='square' # (FontAwesome names without the `fa-` prefix)
)


net = cv2.dnn.readNetFromDarknet('c:\\dog2\\yolov4.cfg','c:\\dog2\\yolov4.weights')

classes =['Pomeranian',
'Bichon_frise',
'Maltese',
'Welsh_corgi',
'Poodle',
'Border_collie',
'Dachshund',
'Chihuahua',
'French_bulldog',
'Shih_tzu',
'Siberian_husky',
'Jindo']

font = cv2.FONT_HERSHEY_PLAIN
colors = np.random.uniform(0, 255, size=(100, 3))

# Display function
# ================
def view(button):
    cap = cv2.VideoCapture(0)
    display_handle=display(None, display_id=True)
    i = 0
    while True:
        _, frame = cap.read()
        frame2 = cv2.flip(frame, 1) # if your camera reverses your image
        
        
        #print(frame.shape)  # 여기까지는 numpy array 였음
        # _, frame3 = cv2.imencode('.jpeg', frame2)
        #print(frame3.shape)  # 여기서 리스트로 변경됨
        #display_handle.update(Image(data=frame3.tobytes()))
        
        
        if stopButton.value==True:   # stop버튼
            cap.release()
            display_handle.update(None)
        
        # 영상의 넓이와 높이를 받음
        width,height = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
        #print(type(frame))
        # 데이터 전처리
        blob = cv2.dnn.blobFromImage(frame2, 1/255, (416, 416), (0,0,0), swapRB=True, crop=False)
        
        # 다크넷 신경망에 넣음
        net.setInput(blob)
        output_layers_names = net.getUnconnectedOutLayersNames()
        layerOutputs = net.forward(output_layers_names)

        boxes = []
        confidences = []
        class_ids = []

        for output in layerOutputs:
            for detection in output:
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                if confidence > 0.8:
                    center_x = int(detection[0]*width)
                    center_y = int(detection[1]*height)
                    w = int(detection[2]*width)
                    h = int(detection[3]*height)

                    x = int(center_x - w/2)
                    y = int(center_y - h/2)

                    boxes.append([x, y, w, h])
                    confidences.append((float(confidence)))
                    class_ids.append(class_id)

        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.8, 0.8)

        if len(indexes)>0:
            for i in indexes.flatten():
                x, y, w, h = boxes[i]
                label = str(classes[class_ids[i]])
                confidence = str(round(confidences[i],2))
                color = colors[i]
                cv2.rectangle(frame2, (x,y), (x+w, y+h), color, 2)
                cv2.putText(frame2, label + " " + confidence, (x, y+20), font, 2, (255,255,255), 2)

        #cv2.imshow('Image', frame)

        #cap = cv2.VideoCapture(input_names[0])
        #win_name = 'Video detection'
        #cv2.namedWindow(win_name) 
        
        # numpy array형식의 frame을 리스트로 변환
        _, frame3 = cv2.imencode('.jpeg', frame2)
        
        # 주피터 노트북에서 실시간으로 display
        display_handle.update(Image(data=frame3.tobytes()))
        
        # 취소키 누르면 while loop문 종료
        key = cv2.waitKey(1)
        if key==27:
            break

        
    out.release()        
    cv2.destroyAllWindows()

            
# Run
# ================
display(stopButton)
thread = threading.Thread(target=view, args=(stopButton,))
thread.start()

NMS와 CONFIDENCE를 0.8정도로 높여줘야 사람얼굴 detection하지 않음.

반응형