Source code for pyEdgeEval.preprocess.nms.fast_nms

#!/usr/bin/env python3

import os
import cv2
import numpy as np
from scipy.io import loadmat

from pyEdgeEval._lib import nms

from .toolbox import conv_tri, grad2


# NOTE:
#    In NMS, `if edge < interp: out = 0`, I found that sometimes edge is very close to interp.
#    `edge = 10e-8` and `interp = 11e-8` in C, while `edge = 10e-8` and `interp = 9e-8` in python.
#    ** Such slight differences (11e-8 - 9e-8 = 2e-8) in precision **
#    ** would lead to very different results (`out = 0` in C and `out = edge` in python). **
#    Sadly, C implementation is not expected but needed :(


[docs]def fast_nms( img: np.ndarray, r: int = 1, s: int = 5, m: float = 1.01, half_prec: bool = False, return_as_uint8: bool = False, ) -> np.ndarray: """NMS for binary edges Args: img (np.ndarray): edge (np.uint8 or float) r (int): radius for nms supr s (int): radius for supr boundaries m (float): multiplier for conservative supr Returns: supressed edge References: - https://github.com/pdollar/edges/blob/master/private/edgesNmsMex.cpp Current runtime is around 20ms """ if img.dtype == np.uint8: # NOTE: input image must be normalized between 0~1 img = (img / 255.0).astype(np.float64) assert (img.dtype == np.float64) or ( img.dtype == np.float32 ), f"ERR: input dtype should be float64 or float32 but got {img.dtype}" edge = conv_tri(img, 1) if half_prec: # although there aren't much speed ups from using half precision, # using this may save memory edge = np.float32(edge) ox, oy = grad2(conv_tri(edge, 4)) oxx, _ = grad2(ox) oxy, oyy = grad2(oy) # sometimes oxx + 1e-5 = 0, and causes true divide warnings oxx[np.where(oxx == 0)] = 1e-5 val = oyy * np.sign(-oxy) / oxx ori = np.mod(np.arctan(val), np.pi) # r, s, m = 1, 5, float(1.01) out = nms(edge, ori, r=r, s=s, m=m) if return_as_uint8: out = np.clip(out, 0, 1) # NOTE: in MATLAB, uint8(x) means round(x).astype(uint8) in numpy out = np.round(out * 255).astype(np.uint8) return out
""" Helper Functions: Probably won't use... """
[docs]def nms_process_one_image(img, save_path=None, save=True): """NMS single image Args: img (np.ndarray): edge, model output save_path (str): save path save (bool): if True, save .png Returns: edge """ if save and save_path is not None: assert os.path.splitext(save_path)[-1] == ".png" out = fast_nms(img, r=1, s=5, m=1.01, return_as_uint8=True) if save: # remove cv2 dependencies cv2.imwrite(save_path, out) return out
[docs]def nms_results(result_dir, save_dir, loader): assert os.path.exists(result_dir), f"ERR: {result_dir} does not exist" if not os.path.isdir(save_dir): os.makedirs(save_dir) for file in os.listdir(result_dir): save_name = os.path.join(save_dir, f"{os.path.splitext(file[0])}.png") if os.path.isfile(save_name): print(f"file: {save_name} exists... skipping") img_path = os.path.join(result_dir, file) # load image img = loader(img_path) # nms preprocess out = fast_nms(img, r=1, s=5, m=1.01, return_as_uint8=True) # save output cv2.imwrite(save_name, out)
[docs]def nms_all_results( model_name_list, result_dir, save_dir, key=None, file_format=".mat" ): if not isinstance(model_name_list, list): model_name_list = [model_name_list] assert file_format in {".mat", ".npy"} assert os.path.isdir(result_dir) for model_name in model_name_list: model_save_dir = os.path.join(save_dir, model_name) if not os.path.isdir(model_save_dir): os.makedirs(model_save_dir) for file in os.listdir(result_dir): save_name = os.path.join( model_save_dir, "{}.png".format(os.path.splitext(file)[0]) ) if os.path.isfile(save_name): continue if os.path.splitext(file)[-1] != file_format: continue abs_path = os.path.join(result_dir, file) if file_format == ".mat": assert key is not None image = loadmat(abs_path)[key] elif file_format == ".npy": image = np.load(abs_path) else: raise NotImplementedError out = fast_nms(image, r=1, s=5, m=1.01, return_as_uint8=True) # save output cv2.imwrite(save_name, out)