{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Salient Object Detection (SOD) - how to remove background from whales images\n\nSalient Object Detection (SOD) aims at segmenting the most visually attractive objects in an image. It is widely\nused in many fields, such as visual tracking and image segmentation. Recently, with the development of deep convolutional neural networks (CNNs), especially the rise of Fully Convolutional Networks (FCN) in image segmentation, the salient object detection has been improved significantly. (Source: https://arxiv.org/pdf/2005.09007.pdf)\n\nDuring my reearch I found SOD survey [RGB-D Salient Object Detection: A Survey](https://github.com/taozh2017/RGBD-SODsurvey). This is a survey to review related RGB-D SOD models along with benchmark datasets, and provide a comprehensive evaluation for these models. Authors also collect related review papers for SOD and light field SOD models.\n\n<div align=\"center\"><img src=\"https://github.com/taozh2017/RGBD-SODsurvey/raw/master/figures/Fig0.jpg\" width=640/></div>\nIn this experiment I use U2-Net: U Square Net whis is described in <a href=\"https://arxiv.org/pdf/2005.09007.pdf\">U2-Net: Going Deeper with Nested U-Structure for Salient Object Detection</a>\n<p>&nbsp;</p>\nU2-Net Architecture:\n<div align=\"center\"><img src=\"https://github.com/xuebinqin/U-2-Net/raw/master/figures/U2NETPR.png\" width=480/></div>\n\nI use this for solution with LoFTR features matching which you can see in my notebook [Whales feature matching LoFTR - Kornia](https://www.kaggle.com/remekkinas/whales-feature-matching-loftr-kornia) \n    \n<div align=\"center\"><img src=\"https://i.ibb.co/yQzhzb8/Lo-FTR-BR1.jpg\"/></div>\n\n<div class=\"alert alert-warning\">Note: My goal was to implement and share tool for experimentations  - I was not creating full dataset using this. This is part of your journey. Enjoy experimenting and progressing!</div>","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"markdown","source":"## 1. SETUP NOTEBOOK (MODULES)\nLet's clone U-2-Net repository","metadata":{}},{"cell_type":"code","source":"!git clone https://github.com/shreyas-bk/U-2-Net\n    \nimport sys\nsys.path.append('./U-2-Net')","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:17.781547Z","iopub.execute_input":"2022-03-07T14:47:17.782515Z","iopub.status.idle":"2022-03-07T14:47:19.639523Z","shell.execute_reply.started":"2022-03-07T14:47:17.782382Z","shell.execute_reply":"2022-03-07T14:47:19.638398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\n\nfrom data_loader import RescaleT\nfrom data_loader import ToTensor\nfrom data_loader import ToTensorLab\nfrom data_loader import SalObjDataset\nfrom model import U2NET \nfrom model import U2NETP \n\nfrom IPython.display import display\nfrom PIL import Image as Img\n\nfrom torchvision import transforms\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.autograd import Variable","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:19.641814Z","iopub.execute_input":"2022-03-07T14:47:19.642114Z","iopub.status.idle":"2022-03-07T14:47:23.074524Z","shell.execute_reply.started":"2022-03-07T14:47:19.642079Z","shell.execute_reply":"2022-03-07T14:47:23.073566Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. TAKE CANDIDATES FOR BACKGROUND REMOVAL","metadata":{}},{"cell_type":"markdown","source":"For this demo I use cropped images from dataset provided by @phalanx [cropped&resized(512x512) dataset using detic](https://www.kaggle.com/c/happy-whale-and-dolphin/discussion/305503). Thank you for contributing in this competition.","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_csv(\"../input/happy-whale-and-dolphin/train.csv\")\n\ninput_path = \"../input/whale2-cropped-dataset/cropped_train_images/cropped_train_images\"","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:23.075913Z","iopub.execute_input":"2022-03-07T14:47:23.076210Z","iopub.status.idle":"2022-03-07T14:47:23.174241Z","shell.execute_reply.started":"2022-03-07T14:47:23.076170Z","shell.execute_reply":"2022-03-07T14:47:23.173136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's take one of TOP10 individual. You can experiment with other individuals. Quality of prediction depends on photo but I will work on improving prediction (I will probably train this model on custom data).","metadata":{}},{"cell_type":"code","source":"train_df.individual_id.value_counts().head(10)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:23.175731Z","iopub.execute_input":"2022-03-07T14:47:23.176094Z","iopub.status.idle":"2022-03-07T14:47:23.207628Z","shell.execute_reply.started":"2022-03-07T14:47:23.176017Z","shell.execute_reply":"2022-03-07T14:47:23.206767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"For experiments we use only one individual ID=281504409737... just to check solution performance and quality of mask.","metadata":{}},{"cell_type":"code","source":"img_to_draw = [ input_path + '/' + file for file in train_df.query(\"individual_id == '281504409737'\").sample(25).image]","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:23.210162Z","iopub.execute_input":"2022-03-07T14:47:23.210979Z","iopub.status.idle":"2022-03-07T14:47:23.224713Z","shell.execute_reply.started":"2022-03-07T14:47:23.210927Z","shell.execute_reply":"2022-03-07T14:47:23.223853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axes = plt.subplots(5, 5, figsize=(20,20))\n\nfor idx, img in enumerate(img_to_draw):\n    i = idx % 5 \n    j = idx // 5\n    image = Img.open(img)\n    iar_shp = np.array(image).shape\n    axes[i, j].axis('off')\n    axes[i, j].imshow(image)\n    \nplt.subplots_adjust(wspace=0.05, hspace=0.05)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:23.225807Z","iopub.execute_input":"2022-03-07T14:47:23.226529Z","iopub.status.idle":"2022-03-07T14:47:26.320380Z","shell.execute_reply.started":"2022-03-07T14:47:23.226495Z","shell.execute_reply":"2022-03-07T14:47:26.319267Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3. SOD - U2-Net PREDICTION","metadata":{}},{"cell_type":"code","source":"THRESHOLD = 0.3","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:26.321945Z","iopub.execute_input":"2022-03-07T14:47:26.323129Z","iopub.status.idle":"2022-03-07T14:47:26.327520Z","shell.execute_reply.started":"2022-03-07T14:47:26.323069Z","shell.execute_reply":"2022-03-07T14:47:26.326358Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def normPRED(d):\n    ma = torch.max(d)\n    mi = torch.min(d)\n    dn = (d-mi)/(ma-mi)\n    return dn\n\n\ndef pred_unet(model, imgs):\n    \n    test_salobj_dataset = SalObjDataset(img_name_list = imgs, lbl_name_list = [], transform = transforms.Compose([RescaleT(320),ToTensorLab(flag=0)]))\n    test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1, shuffle=False, num_workers = 1)\n    \n    for i_test, data_test in enumerate(test_salobj_dataloader):\n        \n        inputs_test = data_test['image']\n        inputs_test = inputs_test.type(torch.FloatTensor)\n\n        if torch.cuda.is_available():\n            inputs_test = Variable(inputs_test.cuda())\n        else:\n            inputs_test = Variable(inputs_test)\n\n        d1, d2, d3, d4, d5, d6, d7 = net(inputs_test)\n\n        predict = d5[:,0,:,:]\n        predict = normPRED(predict)\n        \n        del d1, d2, d3, d4, d5, d6, d7\n\n        predict = predict.squeeze()\n        predict_np = predict.cpu().data.numpy()\n\n        # Masked image - using threshold you can soften/sharpen mask boundaries\n        predict_np[predict_np > THRESHOLD] = 1\n        predict_np[predict_np <= THRESHOLD] = 0\n        mask = Img.fromarray(predict_np*255).convert('RGB')\n        image = Img.open(imgs[0])\n        imask = mask.resize((image.width, image.height), resample=Img.BILINEAR)\n        back = Img.new(\"RGB\", (image.width, image.height), (255, 255, 255))\n        mask = imask.convert('L')\n        im_out = Img.composite(image, back, mask)\n        \n        # Sailient mask \n        salient_mask = np.array(image)\n        mask_layer = np.array(imask)        \n        mask_layer[mask_layer == 255] = 50 # offest on RED channel\n        salient_mask[:,:,0] += mask_layer[:,:, 0]\n        salient_mask = np.clip(salient_mask, 0, 255) \n    \n    return np.array(im_out), np.array(image), np.array(salient_mask), np.array(mask)","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:26.328834Z","iopub.execute_input":"2022-03-07T14:47:26.329415Z","iopub.status.idle":"2022-03-07T14:47:26.347485Z","shell.execute_reply.started":"2022-03-07T14:47:26.329373Z","shell.execute_reply":"2022-03-07T14:47:26.346523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"UNET2_SMALL = False","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:26.348928Z","iopub.execute_input":"2022-03-07T14:47:26.349428Z","iopub.status.idle":"2022-03-07T14:47:26.363111Z","shell.execute_reply.started":"2022-03-07T14:47:26.349385Z","shell.execute_reply":"2022-03-07T14:47:26.362065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%capture\n\nif UNET2_SMALL:\n    model_dir = \"./U-2-Net/u2netp.pth\"  # Faster ... a lot (!) but less accurate\n    net = U2NETP(3,1) \nelse:\n    model_dir = \"../input/u-square-net-model/u2net.pth\"\n    net = U2NET(3,1) \n\n\nif torch.cuda.is_available():\n    net.load_state_dict(torch.load(model_dir))\n    net.cuda()\nelse:        \n    net.load_state_dict(torch.load(model_dir, map_location=torch.device('cpu')))\n\nnet.eval()","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:26.364613Z","iopub.execute_input":"2022-03-07T14:47:26.365195Z","iopub.status.idle":"2022-03-07T14:47:29.358005Z","shell.execute_reply.started":"2022-03-07T14:47:26.365150Z","shell.execute_reply":"2022-03-07T14:47:29.357252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. U2-Net RESULT VISUALIZATION","metadata":{}},{"cell_type":"code","source":"fig, axes = plt.subplots(25, 1, figsize=(600,200))\n\nfor idx, img in enumerate(img_to_draw):\n    image, im_oryg, sal_map, mask = pred_unet(net, [img_to_draw[idx]]) \n    result = np.concatenate((im_oryg, sal_map, image), axis=1)\n    result_img = Img.fromarray(result)\n    axes[idx].axis('off')\n    axes[idx].imshow(result_img)\n\nplt.subplots_adjust(wspace=0.05, hspace=0)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:47:29.359229Z","iopub.execute_input":"2022-03-07T14:47:29.359548Z","iopub.status.idle":"2022-03-07T14:48:46.386159Z","shell.execute_reply.started":"2022-03-07T14:47:29.359507Z","shell.execute_reply":"2022-03-07T14:48:46.384799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<div class=\"alert alert-success\" role=\"alert\">\n    I really appreciate if you suport my work. <b>Voting is more then welcome. It motivates me a lot</b> for sharing part of solution/tools you can use in this competition.\n    \nMy other work in this competition:\n    <ul>\n        <li> <a href=\"https://www.kaggle.com/remekkinas/whales-feature-matching-loftr-kornia\">Whales feature matching LoFTR - Kornia</a></li>\n    </ul>\n    \n</div>","metadata":{}},{"cell_type":"markdown","source":"## 5. IDEAS FOR IMPROVEMENT\n\n### A. REMOVE WHITE AREA OUTSIDE SAILENT MAP - PROTOTYPE\nLet's find bbox for Sailent Map and remove white area outside it. This is just prototype - some code refactoring is needed. ","metadata":{}},{"cell_type":"code","source":"fig, axes = plt.subplots(4, 3, figsize=(30,40))\n\nfor idx, img in enumerate(img_to_draw):\n    \n    if idx > 3:\n        break\n        \n    image, im_oryg, sal_map, mask = pred_unet(net, [img_to_draw[idx]])\n    \n    ymin = np.nonzero((mask[:] != 0).argmax(axis = 1))[0][0]\n    ymax = np.nonzero((mask[:] != 0).argmax(axis = 1))[0][-1]\n    xmin = np.nonzero((mask[:] != 0).argmax(axis = 0))[0][0]\n    xmax = np.nonzero((mask[:] != 0).argmax(axis = 0))[0][-1]\n    \n    img = cv2.rectangle(image.copy(), (xmin, ymin), (xmax, ymax), (255,0,0), 2)\n    \n    crop_img = image[ymin:ymax, xmin:xmax]\n    crop_img = cv2.resize(crop_img, (image.shape[0], image.shape[1]), interpolation = cv2.INTER_AREA)\n    \n    result = np.concatenate((img, crop_img), axis=1)\n    result_img = Img.fromarray(result)\n    \n    axes[idx, 0].imshow(Img.fromarray(sal_map))\n    axes[idx, 1].imshow(Img.fromarray(img))\n    axes[idx, 2].imshow(Img.fromarray(crop_img))\n\nplt.subplots_adjust(wspace=0.1, hspace=0)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-03-07T14:48:46.388678Z","iopub.execute_input":"2022-03-07T14:48:46.389330Z","iopub.status.idle":"2022-03-07T14:49:04.069858Z","shell.execute_reply.started":"2022-03-07T14:48:46.389270Z","shell.execute_reply":"2022-03-07T14:49:04.068791Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This notebook needs some improvement:\n- code refactoring - I decided to share idea as fast as possible (I have done many research to find the best SOD solution)\n- model training on custom data to improve object detection prediction","metadata":{}}]}