{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Jigsaw puzzle\nAccording to the [GWHD paper](https://arxiv.org/pdf/2005.02162.pdf) the images were split into 1024x1024 squared patches containing roughly 20 to 60 heads with only a few heads crossing the edges.\n![gwhd_1](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F864684%2F9068f7f8dd526caa3657407c477e418f%2Fgwhd_splits.jpg?generation=1589069353753211&alt=media)\n\nIn this notebook you can find a raw code to solve this jigsaw puzzle.\n\n## How it works\n- From all of the images, it saves the left, right, top, and bottom edges. (1x2 has only left-right edges)\n- In all possible combinations, it calculates a cosine distance (simple brute force algorithm)\n- In the case of 1x2, it saves the top n matches.\n- In the case of 2x2, it calculates a distance-sum on all possible (and valid) combination\n- Store the top n matches\n- Note: It generates only for one source per run. You have to modify the `TARGET` config (see below) and re-run everything.\n\n## Limitations:\n- The images from the `arvalis_3` source have very similar edges (mostly brown dirt), and this script won't work. (If you have any idea how to solve this, please leave me a comment)\n- The `arvalis_1` and the `rres_1` sources have 2x3 patches (see the image above), this code is not working with 2x3 patches yet\n- There are a few issues with the bounding boxes around edges. Ideally, we should merge the split boxes.\n\n## Notations:\n- **cX** columns in the generated files mean cell-X (for example 2x2: c1=top left, c2=bottom left, c3=top-right, c4=bottom-right). The values in theses columns are the original image id-s.\n\n## TODO:\n- 2x3 solution\n- Merging the split boxes\n- Create box dataframe with the same format we already have\n- Fixing `arvalis_3` bug\n- Create a dataset\n"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport os\nimport ast\nimport hashlib\nimport matplotlib.pyplot as plt\n\nfrom scipy.spatial.distance import cosine\n%matplotlib inline\n\n# Do not modify these!\nTARGETS = ['arvalis_1', 'arvalis_2', 'arvalis_3', 'inrae_1', 'usask_1', 'rres_1', 'ethz_1']\nN_PATCH = [(2, 3), (2, 2), (2, 2), (2, 2), (1, 2), (2, 3), (1, 2)]\nTHRESHOLDS = [0.0, 0.12, 0.0, 0.1, 0.12, 0.0, 0.12]\nN_FULL_SIZE_IMAGES = [239, 51, 152, 44, 100, 72, 375]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ======================================\n# Set these values for your environment\n# ======================================\nDIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'\n\n# Target source index (see above: TARGETS)\n# Currently working: [1, *2, 3, 4, 6]\n# *Image similarity issue with arvalis_3 (idx: 2)\nTARGET = 3\n\n# On Kaggle the script won't save images\nIS_KAGGLE_ENV = True\n\n# DIR_TARGET_ROOT = f'{DIR_INPUT}/full_images'\nDIR_TARGET_ROOT = './'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f\"{DIR_INPUT}/train.csv\")\ntrain_df = train_df[train_df['source'] == TARGETS[TARGET]]\ntrain_df.reset_index(drop=True, inplace=True)\n\nimage_ids = train_df['image_id'].unique()\n\n# From Andrew's kernel\ntrain_df[['x', 'y', 'w', 'h']] = pd.DataFrame(\n    np.stack(train_df['bbox'].apply(lambda x: ast.literal_eval(x)))).astype(np.float32)\ntrain_df['x1'] = train_df['x'] + train_df['w']\ntrain_df['y1'] = train_df['y'] + train_df['h']\n\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Calculate edge similarity"},{"metadata":{"trusted":true},"cell_type":"code","source":"te = []  # top edges\nbe = []  # bottom edges\nle = []  # left edges\nre = []  # right edges\n\nids = [] # image ids\n\nres_rl = [] # Results right-left\nres_bt = [] # Results bottom-top\n\nfor image_id in image_ids:\n    image = cv2.imread(DIR_TRAIN + '/{}.jpg'.format(image_id))\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    \n    le.append(image[:, :1, :].flatten() / 255.0)\n    re.append(image[:, -1:, :].flatten() / 255.0)\n\n    # number of patch-rows\n    if N_PATCH[TARGET][0] > 1:\n        te.append(image[:1, :, :].flatten() / 255.0)\n        be.append(image[-1:, :, :].flatten() / 255.0)\n\n    ids.append(image_id)\n\n# left-right edge distances\nfor i, r in enumerate(re):\n    for j, l in enumerate(le):\n        if i == j:\n            continue\n\n        res_rl.append({\n            'right_edge': i,\n            'left_edge': j,\n            'cd': cosine(r, l)\n        })\n\n# number of patch-rows\nif N_PATCH[TARGET][0] > 1:\n\n    # bottom-top edge distances\n    for i, b in enumerate(be):\n        for j, t in enumerate(te):\n            if i == j:\n                continue\n\n            res_bt.append({\n                'bottom_edge': i,\n                'top_edge': j,\n                'cd': cosine(b, t)\n            })","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"res_rl = pd.DataFrame(res_rl)\nres_rl.sort_values(by='cd', inplace=True)\nres_rl = res_rl[res_rl['cd'] < THRESHOLDS[TARGET]]\n\n# number of patch-rows\nif N_PATCH[TARGET][0] > 1:\n    res_bt = pd.DataFrame(res_bt)\n    res_bt.sort_values(by='cd', inplace=True)\n    res_bt = res_bt[res_bt['cd'] < THRESHOLDS[TARGET]]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Calculate distance sum"},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_results(data, n):\n    sum_dist = pd.DataFrame(data)\n    sum_dist.sort_values(by='dist', inplace=True)\n    sum_dist = sum_dist.iloc[0:N_FULL_SIZE_IMAGES[TARGET]]\n    sum_dist.reset_index(drop=True, inplace=True)\n\n    result = sum_dist.copy()\n    result.drop(columns=['dist'], inplace=True)\n    result['image_id'] = ''\n\n    for i, row in enumerate(sum_dist.iterrows()):\n        row = row[1]\n\n        for c in range(n):\n            result.loc[i, f'c{c+1}'] = ids[int(row[f'c{c+1}'])]\n\n        newid = hashlib.md5(f'full_image_{TARGETS[TARGET]}_{i}'.encode()).hexdigest()[:10]\n        result.loc[i, 'image_id'] = newid\n\n    result.to_csv(f'{DIR_TARGET_ROOT}/full_size_{TARGETS[TARGET]}.csv', index=False)\n    return result\n\n\nif N_PATCH[TARGET] == (1, 2):\n    sum_dist = []\n    \n    res_rl = res_rl.iloc[0:N_FULL_SIZE_IMAGES[TARGET]]\n    res_rl.reset_index(drop=True, inplace=True)\n\n    for i, row in res_rl.iterrows():\n        sum_dist.append({\n            'c1': row['right_edge'],\n            'c2': row['left_edge'],\n            'dist': row['cd']\n        })\n        \n    result = save_results(sum_dist, n = N_PATCH[TARGET][0] * N_PATCH[TARGET][1])\n    \n    \nelif N_PATCH[TARGET] == (2, 2):\n\n    sum_dist = []\n\n    for i1, r1 in res_rl.iterrows():\n        for i2, r2 in res_rl.iterrows():\n            a = res_bt[(res_bt['bottom_edge'] == r1['right_edge']) & (res_bt['top_edge'] == r2['right_edge'])]\n\n            if a.shape[0] == 0:\n                continue\n\n            b = res_bt[(res_bt['bottom_edge'] == r1['left_edge']) & (res_bt['top_edge'] == r2['left_edge'])]\n\n            if b.shape[0] == 0:\n                continue\n\n\n            sum_dist.append({\n                'c1': r1['right_edge'],\n                'c2': r2['right_edge'],\n                'c3': r1['left_edge'],\n                'c4': r2['left_edge'],\n                'dist': r1['cd'] + r2['cd'] + a['cd'].values[0] + b['cd'].values[0]\n            })\n\n    result = save_results(sum_dist, n = N_PATCH[TARGET][0] * N_PATCH[TARGET][1])\n\n    \nelse:\n    raise NotImplementedError\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Saving results"},{"metadata":{"trusted":true},"cell_type":"code","source":"result.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_image(img):\n    image1 = cv2.imread(DIR_TRAIN + '/{}.jpg'.format(img))\n    return cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)\n\ndef save_image(img, image_id, folder):\n    os.makedirs(folder, exist_ok=True)\n    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)\n    cv2.imwrite(f\"{folder}/{image_id}.jpg\", img)\n\ndef show_image(image):\n    fig, ax = plt.subplots(1, 1, figsize=(16, 16))\n    ax.set_axis_off()\n    ax.imshow(image)    \n    \ndef generate_full_image(patches, rows=2, cols=2):\n\n    i = 0\n    target = np.zeros((rows * 1024, cols * 1024, 3))\n\n    for c in range(cols):\n        for r in range(rows):\n            im = load_image(patches[f'c{i+1}'])\n            target[(1024*r):(1024*(r+1)), (1024*c):(1024*(c+1)), :] = im\n            i += 1\n\n    return target.astype(np.uint8)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not IS_KAGGLE_ENV:\n    for i, row in result.iterrows():\n        image = generate_full_image(row, rows=N_PATCH[TARGET][0], cols=N_PATCH[TARGET][1])\n        save_image(image, row['image_id'], f\"{DIR_TARGET_ROOT}/{TARGETS[TARGET]}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Examples"},{"metadata":{"trusted":true},"cell_type":"code","source":"data = result.iloc[10]\nimage = generate_full_image(data, rows=N_PATCH[TARGET][0], cols=N_PATCH[TARGET][1])\n\nshow_image(image)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Adjust the bounding boxes"},{"metadata":{"trusted":true},"cell_type":"code","source":"COLORS = [(220, 0, 0), (0, 0, 220), (0, 220, 128), (220, 128, 0), (0, 220, 0), (128, 0, 128)]\nOFFSETS = [(0, 0), (0, 1024), (1024, 0), (1024, 1024)]\n\ndef show_full_image_with_boxes(data):\n    boxes = {}\n    image = generate_full_image(data, rows=N_PATCH[TARGET][0], cols=N_PATCH[TARGET][1])\n\n    for c in range(N_PATCH[TARGET][0] * N_PATCH[TARGET][1]):\n\n        # Boxes for cell #c\n        df = train_df[train_df['image_id'] == data[f'c{c+1}']][['x', 'y', 'x1', 'y1']]\n\n        df['x'] = df['x'] + OFFSETS[c][0]\n        df['x1'] = df['x1'] + OFFSETS[c][0]\n\n        df['y'] = df['y'] + OFFSETS[c][1]\n        df['y1'] = df['y1'] + OFFSETS[c][1]\n\n        boxes[f'c{c+1}'] = df\n\n    fig, ax = plt.subplots(1, 1, figsize=(16, 16))\n\n    i = 0\n    for c in range(N_PATCH[TARGET][1]):\n        for r in range(N_PATCH[TARGET][0]):\n            bboxes = boxes[f'c{i+1}'].values.astype(np.uint32)\n\n            for box in bboxes:\n                cv2.rectangle(image, (box[0], box[1]),\n                              (box[2], box[3]), COLORS[i], 2)\n\n            i += 1\n\n    ax.set_axis_off()\n    ax.imshow(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_full_image_with_boxes(result.iloc[10])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_full_image_with_boxes(result.iloc[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_full_image_with_boxes(result.iloc[23])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}