Heooll/dragon-pic-detect

1 2import cv23import imageio4import glob5 6import numpy as np7 8def read_img(name) :9    if name.endswith(".gif"):10        gif = imageio.mimread(name)11        return cv2.cvtColor(gif[0], cv2.COLOR_RGB2BGR)12    else :13        return cv2.cvtColor(imageio.imread(name), cv2.COLOR_RGB2BGR)14 15def image_resize(image, width = None, height = None, inter = cv2.INTER_AREA):16    # initialize the dimensions of the image to be resized and17    # grab the image size18    dim = None19    (h, w) = image.shape[:2]20    # if both the width and height are None, then return the21    # original image22    if width is None and height is None:23        return image24    # check to see if the width is None25    if width is None:26        # calculate the ratio of the height and construct the27        # dimensions28        r = height / float(h)29        dim = (int(w * r), height)30    # otherwise, the height is None31    else:32        # calculate the ratio of the width and construct the33        # dimensions34        r = width / float(w)35        dim = (width, int(h * r))36 37    # resize the image38    resized = cv2.resize(image, dim, interpolation = inter)39    # return the resized image40    return resized41 42def minEnclosingCircleArea(pts) :43    _, r = cv2.minEnclosingCircle(pts)44    return r * r * np.pi45 46class DragonDetector(object) :47    def __init__(self, template_image_pattern = 'template*.png', image_resolutions = [20, 60, 100, 200, 400], match_point_threshold = 5, circle_area_threshold = 0.2) :48        self.templates = [read_img(fn) for fn in glob.glob(template_image_pattern)]49        self.match_point_threshold = match_point_threshold50        self.circle_area_threshold = circle_area_threshold51        self.image_resolutions = image_resolutions52        self.sift = cv2.SIFT_create()53        self.template_sifts = [(self.sift.detectAndCompute(img, None), img.shape) for img in self.templates]54        # FLANN parameters55        FLANN_INDEX_KDTREE = 156        index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)57        search_params = dict(checks=50)   # or pass empty dictionary58        self.flann = cv2.FlannBasedMatcher(index_params,search_params)59        print(f' -- {len(self.templates)} dragon templates loaded')60 61    def is_dragon_impl(self, kps, imgnp) :62        (kp1, des1), template_shape = kps63        kp2, des2 = self.sift.detectAndCompute(imgnp, None)64        try :65            matches = self.flann.knnMatch(des1, des2, k = 2)66        except Exception :67            return False68        good = []69        for m, n in matches :70            if m.distance < 0.7 * n.distance:71                good.append(m)72        if len(good) > self.match_point_threshold :73            src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)74            dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)75            M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)76            h, w, d = template_shape77            pts = np.float32([ [0,0],[0,h-1],[w-1,h-1],[w-1,0] ]).reshape(-1,1,2)78            try :79                dst = cv2.perspectiveTransform(pts,M)80                if cv2.contourArea(dst) > self.circle_area_threshold * minEnclosingCircleArea(dst) :81                    return True82            except Exception :83                pass84        return False85 86    def is_dragon(self, img_np) :87        found = False88        for template in self.template_sifts :89            for w in self.image_resolutions :90                found = found or self.is_dragon_impl(template, image_resize(img_np, width = w, inter = cv2.INTER_LINEAR))91        return found92 93class DragonDetectorFast(DragonDetector) :94    def __init__(self) :95        super().__init__('template1.png', [200])96 97def test() :98    import tqdm99    tp, tn, fp, fn = 0, 0, 0, 0100    det = DragonDetectorFast()101    for img in tqdm.tqdm(glob.glob("cascade - Copy\\n\\*")) :102        img2 = read_img(img)103        found = det.is_dragon(img2)104        if found :105            fp += 1106        else :107            tn += 1108 109    for img in tqdm.tqdm(glob.glob("cascade - Copy\\p\\*")) :110        img2 = read_img(img)111        found = det.is_dragon(img2)112        if found :113            tp += 1114        else :115            fn += 1116    print(f'tp = {tp}')117    print(f'tn = {tn}')118    print(f'fp = {fp}')119    print(f'fn = {fn}')120 121if __name__ == '__main__' :122    test()123