Removing grabcut and classic segmentations (#70)

Removing grabcut and classic segmentations, its tests, and now-redundant arguments for U-Net.

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
     ...

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)

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))

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)