{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Split overlapping bounding boxes\n\n**Detailed description:** https://towardsdatascience.com/split-overlapping-bounding-boxes-in-python-e67dc822a285","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np \nimport pandas as pd \nfrom fastai.vision import *\nimport geopandas as gpd\nimport shapely\nfrom shapely.geometry import Polygon, LineString\nfrom tqdm import tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"path = Path('/kaggle/input/global-wheat-detection/')","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"df = pd.read_csv(path/'train.csv')\ndf","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def bbox2mask(x):\n    labels = np.array(x)\n    mask = torch.zeros(1024,1024)\n    for l in labels:\n        mask[l[1]:l[1]+l[3], l[0]:l[0]+l[2]] = 1\n    return mask\n\ndef bbox_center(x):\n    labels = np.array(x)\n    mask = torch.zeros(1024,1024)\n    for l in labels:\n        mask[(2*l[1]+l[3])//2, (2*l[0]+l[2])//2] = 1\n    return mask\n\ndef box2polygon(x):\n    return Polygon([(x[0], x[1]), (x[0]+x[2], x[1]), (x[0]+x[2], x[1]+x[3]), (x[0], x[1]+x[3])])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"boxes = df.groupby('image_id').agg({'bbox' : lambda x : list(x)})\nbox = boxes.iloc[2]\nfile = str(path/'train'/box.name) + '.jpg'\nimg = open_image(file).data.numpy().transpose(1,2,0)\nbbox = np.array([eval(l) for l in box.bbox]).astype(int).tolist()\nmask = bbox2mask(bbox)\ngdf = gpd.GeoDataFrame({'geometry': [box2polygon(b) for b in bbox]})\ngdf.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def slice_box(box_A:Polygon, box_B:Polygon, margin=10, line_mult=10):\n    \"Returns box_A sliced according to the distance to box_B.\"\n    vec_AB = np.array([box_B.centroid.x - box_A.centroid.x, box_B.centroid.y - box_A.centroid.y])\n    vec_ABp = np.array([-(box_B.centroid.y - box_A.centroid.y), box_B.centroid.x - box_A.centroid.x])\n    vec_AB_norm = np.linalg.norm(vec_AB)\n    split_point = box_A.centroid + vec_AB/2 - (vec_AB/vec_AB_norm)*margin\n    line = LineString([split_point-line_mult*vec_ABp, split_point+line_mult*vec_ABp])\n    split_box = shapely.ops.split(box_A, line)\n    if len(split_box) == 1: return split_box, None, line\n    is_center = [s.contains(box_A.centroid) for s in split_box]\n    if sum(is_center) == 0: \n        warnings.warn('Polygon do not contain the center of original box, keeping the first slice.')\n        return split_box[0], None, line\n    where_is_center = np.argwhere(is_center).reshape(-1)[0]\n    where_not_center = np.argwhere(~np.array(is_center)).reshape(-1)[0]\n    split_box_center = split_box[where_is_center]\n    split_box_out = split_box[where_not_center]\n    return split_box_center, split_box_out, line","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"inter = gdf.loc[gdf.intersects(gdf.iloc[20].geometry)]\n\nbox_A = inter.iloc[0].values[0]\nbox_B = inter.iloc[1].values[0]\npolyA, _, lineA = slice_box(box_A, box_B, margin=10, line_mult=1.2)\npolyB, _, lineB = slice_box(box_B, box_A, margin=10, line_mult=1.2)\n\nboxes = gpd.GeoDataFrame({'geometry': [box_A, box_B]})\ncentroids =  gpd.GeoDataFrame({'geometry': [box_A.centroid, box_B.centroid]})\nsplited_boxes = gpd.GeoDataFrame({'geometry': [polyA, polyB]})\nlines = gpd.GeoDataFrame({'geometry': [lineA, lineB]})\n\nfig, ax = plt.subplots(dpi=120)\nboxes.plot(ax=ax, facecolor='gray', edgecolor='k', alpha=0.5)\ncentroids.plot(ax=ax, c='k')\nax.axis('off');\n\nfig, ax = plt.subplots(dpi=120)\nboxes.plot(ax=ax, facecolor='gray', edgecolor='k', alpha=0.1)\nsplited_boxes.plot(ax=ax, facecolor='olive', edgecolor='k')\ncentroids.plot(ax=ax, c='k')\nlines.plot(ax=ax, color='k')\nax.axis('off');","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def intersection_list(polylist):\n    r = polylist[0]\n    for p in polylist:\n        r = r.intersection(p)\n    return r\n    \ndef slice_one(gdf, index):\n    inter = gdf.loc[gdf.intersects(gdf.iloc[index].geometry)]\n    if len(inter) == 1: return inter.geometry.values[0]\n    box_A = inter.loc[index].values[0]\n    inter = inter.drop(index, axis=0)\n    polys = []\n    for i in range(len(inter)):\n        box_B = inter.iloc[i].values[0]\n        polyA, *_ = slice_box(box_A, box_B)\n        polys.append(polyA)\n    return intersection_list(polys)\n\ndef slice_all(gdf):\n    polys = []\n    for i in range(len(gdf)):\n        polys.append(slice_one(gdf, i))\n    return gpd.GeoDataFrame({'geometry': polys})","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"res_df = slice_all(gdf)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(10,5), dpi=120)\ngdf.plot(ax=ax1, alpha=0.5, color='gray')\n#gdf.plot(ax=ax2, alpha=0.1, facecolor='gray')\nres_df.plot(ax=ax2, alpha=0.5, color='olive')\nax1.axis('equal')\nax2.axis('equal')\nax1.set_title('Original bounding boxes')\nax2.set_title('Splited bounding boxes')\nfig.tight_layout()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Rasterize polygons","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import rasterio.features","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"raster = rasterio.features.rasterize(res_df.geometry, out_shape=(1024,1024), merge_alg=rasterio.enums.MergeAlg.replace)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, axes = plt.subplots(ncols=2, dpi=120)\naxes[0].imshow(img)\naxes[0].imshow(mask, alpha=0.4)\naxes[1].imshow(img)\naxes[1].imshow(raster, alpha=0.4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Save masks","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import PIL\nimport zipfile\nimport cv2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"mask = cv2.imencode('.png', (raster*255).astype(np.uint8))[1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"boxes = df.groupby('image_id').agg({'bbox' : lambda x : list(x)})\n\nwith zipfile.ZipFile('split_masks.zip', 'w') as mask_out:\n    for i in progress_bar(range(len(boxes))):\n        box = boxes.iloc[i]\n        file = str(path/'train'/box.name) + '.jpg'\n        img = open_image(file).data.numpy().transpose(1,2,0)\n        bbox = np.array([eval(l) for l in box.bbox]).astype(int).tolist()\n        gdf = gpd.GeoDataFrame({'geometry': [box2polygon(b) for b in bbox]})\n        res_df = slice_all(gdf)\n        raster = rasterio.features.rasterize(res_df.geometry, out_shape=(1024,1024), merge_alg=rasterio.enums.MergeAlg.replace)\n        mask = cv2.imencode('.png', (raster*255).astype(np.uint8))[1]\n        mask_out.writestr(f'{box.name}.png', mask)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}