天天看點

Caffe-ssd測試模型代碼

# _*_ coding:UTF-8 _*_
#開發作者   :   ZhangRong z00520111
#開發時間   :   2020/5/13  10:09
#檔案名稱   :   firetest.py
#開發工具   :   PyCharm
#Description:
#Copyright @ Huawei Technologies Co., Ltd. 2019-2020. All rights reserved.
import cv2
import os
import numpy as np
import glob
import json

file = open("result.txt", "a+")
nfile = open("nresult.txt", "a+")
resultpath = "./detection-result/"
nresultpath = "./detection-nresult/"
minScore = 0.8
num = 0


def isexsist(path):
    isExists = os.path.exists(path)

    # 判斷結果
    if not isExists:
        # 如果不存在則建立目錄
        # 建立目錄操作函數
        os.makedirs(path)


fileaddress = []


def walkFile(file):
    for root, dirs, files in os.walk(file):

        # root 表示目前正在通路的檔案夾路徑
        # dirs 表示該檔案夾下的子目錄名list
        # files 表示該檔案夾下的檔案list
        # 周遊檔案
        for f in files:
            fileaddress.append(os.path.join(root, f))
        # # 周遊檔案
        # for f in files:
        #     print(os.path.join(root, f))
        #
        # # 周遊所有的檔案夾
        # for d in dirs:
        #     print(os.path.join(root, d))
    return files, fileaddress


def all_detect(img_path):
    files, fileaddress = walkFile(img_path)
    # print(fileaddress)
    global file
    global nfile
    for index, f in enumerate(fileaddress):
        if str(f).find('.jpg') != -1:
            # isexsist(resultpath)
            # isexsist(nresultpath)
            print(f)
            filepre = f[:f.find('_')]
            # print(f[:f.find('.')])
            # print( files[index][:files[index].find('.')] )
            # file = open(resultpath + files[index][:files[index].find('.')] + ".txt", "a+")
            # nfile = open(nresultpath + "n" +  files[index][:files[index].find('.')]  + ".txt", "a+")
            origimg = cv2.imdecode(np.fromfile(f, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
            resultimage, isshow, rows, cols = detect(origimg, minScore)
            # cv2.namedWindow("resized", 0);
            # cv2.resizeWindow("resized", cols, rows);
            # if isshow == 1:
            #     cv2.imshow("resized", resultimage)
            #     cv2.waitKey(0)


CLASSES = ('background', 'fire','smoke')


def area(left, top, right, bottom):
    return (right - left) * (bottom - top)

def detect(img, minScore):
    isshow = 0
    im_tensor = cv2.dnn.blobFromImage(img, 0.007843, (300, 300), (127.5, 127.5, 127.5), False, False)
    net.setInput(im_tensor)
    rows, cols, chinals = img.shape
    print(rows, cols)
    cvOut = net.forward()
    for detection in cvOut[0, 0, :, :]:
        score = float(detection[2])
        if score > minScore:
            left = int(detection[3] * cols)
            top = int(detection[4] * rows)
            right = int(detection[5] * cols)
            bottom = int(detection[6] * rows)
            weight = int(right - left)
            height = int(bottom - top)
                # cv2.rectangle(img, (int(brandleft), int(brandtop)), (int(brandright), int(brandbottom)), (255, 0, 0), 2)
                # cv2.rectangle(img, (left-int(weight/4), top-int(height/3)-200), (right+int(weight/4), bottom+int(height/3*2)-200), (255, 0, 0), 2)
            cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 2)
            cv2.putText(img, str(CLASSES[int(detection[1])]) + "   score is" + str(score), (left, top + 50), 1, 2,(255, 255, 0), 2)
            print("find fire  ",score, detection[3], detection[4], detection[5], detection[6])
            cv2.namedWindow("result",0)
            cv2.resizeWindow("result",1280,720)
            cv2.imshow("result",img)
            cv2.waitKey(10)
    return img, isshow, rows, cols

def video_detect(path):
    cap=cv2.VideoCapture(path)
    while(1):
        ret,frame=cap.read()
        detect(frame,minScore)
        if  0xFF == ord('q'):
            break

if __name__ == "__main__":
    net = cv2.dnn.readNetFromCaffe("fire_deploy.prototxt",
                                   "models/mobilenet_iter_25000.caffemodel")
    # all_detect("F:/work/測試集/fire/")
    video_detect("F:/work/測試集/火焰/視訊/8.flv")


           

繼續閱讀