diff --git a/README.md b/README.md index a88d3a9..c2062a7 100644 --- a/README.md +++ b/README.md @@ -33,8 +33,6 @@ The following parameters can be used as input arguments for `pipeline.py`: * `-s`, `--stage` : The stage which to run the pipeline until. Options are `'ruler_detection'`, `'binarization'`, and `'measurements'`. Default is `measurement` (running to completion). Running the pipeline and stopping at an earlier stage can be useful for debugging. * `-csv`, `--path_csv` : Path of `.csv` file for the measurement results. (Default is `results.csv`). * `-dpi` : Optional argument to specify resolution of the output image. (Default is `300`.) -* `-g`, `--grabcut` : Use OpenCV's grabcut method in order to improve binarization on blue butterflies. -* `-u`, `--unet` : Use the U-net deep neural network in the binarization step. ## Measurement results @@ -72,9 +70,9 @@ Resulting files: Running the command ``` -$ python pipeline.py -p -u -i ../mothra-data -o ../test_output -csv ../test_output/results.csv +$ python pipeline.py -p -i ../mothra-data -o ../test_output -csv ../test_output/results.csv ``` -in `/mothra` will run the pipeline using the U-net deep neural network on the example data in the folder `/mothra-data`. The file locations should look like this: +in `./mothra` will run the pipeline on the example data in the folder `/mothra-data`. The file locations should look like this: ``` /mothra ... diff --git a/butterfly/binarization.py b/butterfly/binarization.py index b0fd98d..d99be4b 100644 --- a/butterfly/binarization.py +++ b/butterfly/binarization.py @@ -7,7 +7,6 @@ from skimage.morphology import binary_erosion, binary_dilation, selem from skimage.transform import rescale from skimage.util import img_as_bool, img_as_float32, img_as_ubyte -import cv2 as cv from joblib import Memory from fastai.vision import load_learner, Image from pathlib import Path @@ -30,16 +29,6 @@ # Used in find_tags_edge. Percent of width of the image REGION_CUTOFF = 1/3 -# Size of disk used in dilation to get connect butterfly regions -# Used during getting butterfly bounding box in grabcut -DILATION_SIZE = 10 - -# Image downsize percentage used to improve grabcut speed -GRABCUT_RESCALE_FACTOR = 0.25 - -# Number of iterations used in grabcut -GRABCUT_ITERATIONS = 10 - def _convert_image_to_tensor(image): """Auxiliary function. Receives an RGB image and convert it to be processed @@ -128,62 +117,6 @@ def find_tags_edge(image_rgb, top_ruler, axes=None): return label_edge -def grabcut_binarization(bfly_rgb, bfly_bin): - """Extract shape of the butterfly using OpenCV's grabcut. Greatly improves - binarization of blue-colored butterflies. - - Arguments - --------- - bfly_rgb : (M, N, 3) ndarray - Input RGB image of butterfly (ruler and tags cropped out) - bfly_bin : (M, N) ndarray - Binarizaed image of butterfly (ruler and tags cropped out) Expected - to be binarized saturation channel of bfly_rgb. - - Returns - ------- - bfly_grabcut_bin : (M, N) ndarray - Resulting binarized image of butterfly after segmentation by grabcut. - """ - - # Dilation of image to capture butterfly region - selem_arr = selem.disk(DILATION_SIZE) - bfly_bin_dilated = binary_dilation(bfly_bin, selem_arr) - bfly_bin_dilated_markers, _ = ndi.label(bfly_bin_dilated, ndi.generate_binary_structure(2, 1)) - bfly_bin_dilated_regions = regionprops(bfly_bin_dilated_markers) - bfly_bin_dilated_regions_sorted = sorted(bfly_bin_dilated_regions, key=lambda r: r.area, reverse=True) - bfly_region = bfly_bin_dilated_regions_sorted[0] - - # Downscale image to improve grabcut speed - bfly_rgb_rescale = rescale(bfly_rgb, GRABCUT_RESCALE_FACTOR, - multichannel=True) - bfly_rgb_rescale = img_as_ubyte(bfly_rgb_rescale) - - # Determine grabcut highlight region using butterfly region (after rescaling) - padding = 0 - rect = (int(GRABCUT_RESCALE_FACTOR*bfly_region.bbox[1]-padding), - int(GRABCUT_RESCALE_FACTOR*bfly_region.bbox[0]-padding), - int(GRABCUT_RESCALE_FACTOR*bfly_region.bbox[3]+padding), - int(GRABCUT_RESCALE_FACTOR*bfly_region.bbox[2]+padding)) - - # Grabcut - mask = np.zeros(bfly_rgb_rescale.shape[:2], np.uint8) - bgd_model = np.zeros((1,65), np.float64) - fgd_model = np.zeros((1,65), np.float64) - - cv.grabCut(bfly_rgb_rescale, mask, rect, - bgd_model, fgd_model, GRABCUT_ITERATIONS, cv.GC_INIT_WITH_RECT) - - mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8') - bfly_grabcut_rescale = bfly_rgb_rescale*mask2[:, :, np.newaxis] - - # Rescale the image back up and get binary of result - bfly_grabcut = rescale(bfly_grabcut_rescale, bfly_rgb.shape[0]/bfly_rgb_rescale.shape[0]) - bfly_grabcut_bin = np.max(bfly_grabcut, axis=2)>0 - - return bfly_grabcut_bin - - def unet_binarization(bfly_rgb, weights='./models/segmentation.pkl'): """Extract shape of the butterfly using the U-net neural network. @@ -244,7 +177,7 @@ def return_largest_region(img_bin): @memory.cache(ignore=['axes']) -def main(image_rgb, top_ruler, grabcut=False, unet=False, axes=None): +def main(image_rgb, top_ruler, axes=None): """Binarizes and crops properly image_rgb Arguments @@ -268,15 +201,8 @@ def main(image_rgb, top_ruler, grabcut=False, unet=False, axes=None): bfly_rgb = image_rgb[:top_ruler, :label_edge] - if unet: - bfly_bin = unet_binarization(bfly_rgb, weights='./models/segmentation.pkl') - else: - bfly_hsv = color.rgb2hsv(bfly_rgb)[:, :, 1] - rescaled = rescale_intensity(bfly_hsv, out_range=(0, 255)) - thresh_hsv = threshold_otsu(rescaled) - bfly_bin = rescaled > thresh_hsv - if grabcut: - bfly_bin = grabcut_binarization(bfly_rgb, bfly_bin) + # binarizing the input image using U-Net. + bfly_bin = unet_binarization(bfly_rgb, weights='./models/segmentation.pkl') # if the binary image has more than one region, returns the largest one. bfly_bin = return_largest_region(bfly_bin) diff --git a/butterfly/tests/test_binarization.py b/butterfly/tests/test_binarization.py index 42b62fd..b1a5840 100644 --- a/butterfly/tests/test_binarization.py +++ b/butterfly/tests/test_binarization.py @@ -76,24 +76,6 @@ def test_missing_tags(fake_butterfly_no_tags): assert (result >= 399) -def test_grabcut(fake_butterfly_layout): - # mostly as test_main, but testing grabcut method does not interfere with - # expected behavior - butterfly, (rows, cols) = fake_butterfly_layout - picture_2d = butterfly.astype(np.uint8) - picture_3d = np.dstack((picture_2d, - 1/2 * picture_2d, - 1/4 * picture_2d)) # fake RGB image - result = binarization.main(picture_3d, 230, grabcut=True) - y_where, x_where = np.where(result) - x_where_min, x_where_max = np.min(x_where), np.max(x_where) - y_where_min, y_where_max = np.min(y_where), np.max(y_where) - - # assert the "butterfly" is included in the cropped and binarized result - assert((rows - 10 <= y_where_max - y_where_min <= rows + 10) and - (cols - 10 <= x_where_max - x_where_min <= cols + 10)) - - def test_unet(): """Tests U-net segmentation.""" bfly_rgb = imread('./butterfly/tests/test_files/test_input/BMNHE_1297240_147563_dbf1346219110b416ff10347b45e070b1e0d116d.png.rgb') @@ -102,19 +84,3 @@ def test_unet(): output = binarization.unet_binarization(bfly_rgb, weights='./models/segmentation.pkl') assert np.allclose(img_as_bool(output), img_as_bool(bfly_bin)) - - -def test_main(fake_butterfly_layout): - butterfly, (rows, cols) = fake_butterfly_layout - picture_2d = butterfly.astype(np.uint8) - picture_3d = np.dstack((picture_2d, - 1/2 * picture_2d, - 1/4 * picture_2d)) # fake RGB image - result = binarization.main(picture_3d, 230) - y_where, x_where = np.where(result) - x_where_min, x_where_max = np.min(x_where), np.max(x_where) - y_where_min, y_where_max = np.min(y_where), np.max(y_where) - - # assert the "butterfly" is included in the cropped and binarized result - assert((rows - 5 <= y_where_max - y_where_min <= rows + 5) and - (cols - 5 <= x_where_max - x_where_min <= cols + 5)) diff --git a/pipeline.py b/pipeline.py index a73711e..6c21786 100755 --- a/pipeline.py +++ b/pipeline.py @@ -219,15 +219,6 @@ def main(): help='Path of the resulting csv file', default='results.csv') - # Grabcut - parser.add_argument('-g', '--grabcut', - action='store_true', - help='Use grabcut in binarization step') - - # U-nets - parser.add_argument('-u', '--unet', - action='store_true', - help='Use U-nets in binarization step') args = parser.parse_args() # Initialization @@ -290,7 +281,7 @@ def main(): T_space, top_ruler = ruler_detection.main(image_rgb, axes) elif step == 'binarization': - binary = binarization.main(image_rgb, top_ruler, args.grabcut, args.unet, axes) + binary = binarization.main(image_rgb, top_ruler, axes) elif step == 'measurements': points_interest = tracing.main(binary, axes)