patchmatch.py 文件源码

def match(self, im0, im1, hm0, hm1):
        viz = False
        mask0 = self.BG0.segment(im0)
        mask1 = self.BG1.segment(im1)

        im0 = im0 * (mask0>1e-10).astype('uint8')[:,:,np.newaxis]
        im1 = im1 * (mask1>1e-10).astype('uint8')[:,:,np.newaxis]

        if viz:
            viz0 = np.copy(im0)
            viz1 = np.copy(im1)
        pts14 = []
        for chan in range(14):
            h0 = cv2.resize(hm0[:,:,chan], (ORIG_SIZE, ORIG_SIZE))
            h1 = cv2.resize(hm1[:,:,chan], (ORIG_SIZE, ORIG_SIZE))
            y0, x0 = argmax_2d(h0)
            y1, x1 = argmax_2d(h1)

            target = take_patch(im0, y0, x0, PATCH_SIZE)
            region = take_patch(im1, y1, x1, REGION_SIZE)

            res = cv2.matchTemplate(region, target, cv2.TM_CCOEFF_NORMED)
            _, _, _, top_left = cv2.minMaxLoc(res)
            top_left = top_left[::-1]
            center_in_region = (top_left[0] + PATCH_SIZE, top_left[1] + PATCH_SIZE)
            center_in_im1 = (center_in_region[0] + y1-REGION_SIZE,
                    center_in_region[1] + x1-REGION_SIZE)

            if viz:
                cv2.circle(viz0, (x0,y0), 3, (0,0,255), -1)
                cv2.circle(viz1, tuple(center_in_im1[::-1]), 3, (0,0,255), -1)
            pts14.append([x0, y0, center_in_im1[1],center_in_im1[0]])
        if viz:
            mask0 = cv2.cvtColor(mask0, cv2.COLOR_GRAY2RGB).astype('uint8')
            mask1 = cv2.cvtColor(mask1, cv2.COLOR_GRAY2RGB).astype('uint8')
            viz = np.concatenate((mask0, viz0,viz1, mask1),axis=1)
            cv2.imshow("v", viz)
            cv2.waitKey(1)
        return np.array(pts14)
        return viz, np.array(pts14)

        #rv = np.copy(region)
        #rv[center_in_region[0],center_in_region[1]] = (0,0,255)
        #tv = cv2.resize(target, tuple(region.shape[:2][::-1]))

        #hv = np.zeros((region.shape), dtype='float32')
        #res = res - res.min()
        #res = res / res.max() * 255
        #res = cv2.cvtColor(res, cv2.COLOR_GRAY2RGB)
        #hv[PATCH_SIZE:PATCH_SIZE+res.shape[0],PATCH_SIZE:PATCH_SIZE+res.shape[1],:] = res
        #region = np.concatenate((region, rv, tv, hv), axis=1)
        #cv2.imwrite("patchmatch/region{}.png".format(chan), region)