YYC_YOYO.py

from ctypes import *
import math
import random
import os
import time
import cv2 as cv

def sample(probs):
    s = sum(probs)
    probs = [a/s for a in probs]
    r = random.uniform(0, 1)
    for i in range(len(probs)):
        r = r - probs[i]
        if r <= 0:
            return i
    return len(probs)-1
def c_array(ctype, values):
    arr = (ctype*len(values))()
    arr[:] = values
    return arr

class BOX(Structure):
    _fields_ = [("x", c_float),
                ("y", c_float),
                ("w", c_float),
                ("h", c_float)]
class DETECTION(Structure):
    _fields_ = [("bbox", BOX),
                ("classes", c_int),
                ("prob", POINTER(c_float)),
                ("mask", POINTER(c_float)),
                ("objectness", c_float),
                ("sort_class", c_int)]
class IMAGE(Structure):
    _fields_ = [("w", c_int),
                ("h", c_int),
                ("c", c_int),
                ("data", POINTER(c_float))]
class METADATA(Structure):
    _fields_ = [("classes", c_int),
                ("names", POINTER(c_char_p))]

hasGPU = True
cwd = os.path.dirname(__file__)
os.environ['PATH'] = cwd + ';' + os.environ['PATH']
winGPUdll = os.path.join(cwd, "yolo_cpp_dll.dll")
envKeys = list()
lib = CDLL(winGPUdll, RTLD_GLOBAL)
#__________________________________________________
lib.network_width.argtypes = [c_void_p]
lib.network_width.restype = c_int
lib.network_height.argtypes = [c_void_p]
lib.network_height.restype = c_int

get_network_boxes = lib.get_network_boxes
get_network_boxes.argtypes = [c_void_p, c_int, c_int, c_float, c_float, POINTER(c_int), c_int, POINTER(c_int), c_int]
get_network_boxes.restype = POINTER(DETECTION)

free_detections = lib.free_detections
free_detections.argtypes = [POINTER(DETECTION), c_int]

load_net_custom = lib.load_network_custom
load_net_custom.argtypes = [c_char_p, c_char_p, c_int, c_int]
load_net_custom.restype = c_void_p

do_nms_sort = lib.do_nms_sort
do_nms_sort.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]

free_image = lib.free_image
free_image.argtypes = [IMAGE]

load_meta = lib.get_metadata
lib.get_metadata.argtypes = [c_char_p]
lib.get_metadata.restype = METADATA

load_image = lib.load_image_color
load_image.argtypes = [c_char_p, c_int, c_int]
load_image.restype = IMAGE

predict_image = lib.network_predict_image
predict_image.argtypes = [c_void_p, IMAGE]
predict_image.restype = POINTER(c_float)


def detect(net, meta, image, thresh=.5, hier_thresh=.5, nms=.45, debug= False):
    with open("D:\\tools\\python\\MyProjects\\FindImage\\yyc.data") as metaFH:
        metaContents = metaFH.read()
        import re
        match = re.search("names *= *(.*)$", metaContents, re.IGNORECASE | re.MULTILINE)
        result = match.group(1)
        with open(result) as namesFH:
            namesList = namesFH.read().strip().split("\n")
            altNames = [x.strip() for x in namesList]
    im = load_image(image, 0, 0)
    num = c_int(0)
    pnum = pointer(num)
    predict_image(net, im)
    dets = get_network_boxes(net, im.w, im.h, thresh, hier_thresh, None, 0, pnum, 0)
    num = pnum[0]
    do_nms_sort(dets, num, meta.classes, nms)
    res = []
    for j in range(num):
        for i in range(meta.classes):
            if dets[j].prob[i] > 0:
                b = dets[j].bbox
                if altNames is None:
                    nameTag = meta.names[i]
                else:
                    nameTag = altNames[i]
                res.append((nameTag, dets[j].prob[i], (b.x, b.y, b.w, b.h)))
    res = sorted(res, key=lambda x: -x[1])
    free_image(im)
    free_detections(dets, num)
    return res
netMain = None
metaMain = None
altNames = None
def main():
    global metaMain, netMain, altNames
    netMain = load_net_custom("D:\\tools\\python\\MyProjects\\FindImage\\yyc.cfg".encode("ascii"), "D:\\tools\\python\\MyProjects\\FindImage\\yyc_9500.weights".encode("ascii"), 0, 1)  # batch size = 1
    metaMain = load_meta("D:\\tools\\python\\MyProjects\\FindImage\\yyc.data".encode("ascii"))
    NewTims = time.time()
    while True:
        detections = detect(netMain, metaMain, "test.bmp".encode("ascii"), 0.25)
        for i in detections:
            print("名：%s|相似：%s|x:%s|y:%s|w:%s|h:%s" % (i[0], round(i[1], 2), round(i[2][0], 2), round(i[2][1], 2), round(i[2][2], 2), round(i[2][3], 2)))
        print("耗时：", time.time() - NewTims, "___________________________________________________")
        NewTims = time.time()
    # detections = detect(netMain, metaMain, "data/eagle.jpg".encode("ascii"), 0.25)
    # print (detections)


if __name__ == "__main__":
    main()