{"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":"#

SETUP

","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-04-02T19:54:11.619183Z","iopub.execute_input":"2022-04-02T19:54:11.619472Z","iopub.status.idle":"2022-04-02T19:54:12.277667Z","shell.execute_reply.started":"2022-04-02T19:54:11.619422Z","shell.execute_reply":"2022-04-02T19:54:12.276878Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

IMPORT

","metadata":{}},{"cell_type":"code","source":"import os\nimport cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\nimport json\nfrom datetime import datetime\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-04-02T19:54:12.279852Z","iopub.execute_input":"2022-04-02T19:54:12.280164Z","iopub.status.idle":"2022-04-02T19:54:12.286609Z","shell.execute_reply.started":"2022-04-02T19:54:12.280104Z","shell.execute_reply":"2022-04-02T19:54:12.285797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

CONFIGURATION

","metadata":{}},{"cell_type":"code","source":"class CFG:\n # U2-Net\n THRESHOLD=0.5\n UNET2_SMALL = False","metadata":{"execution":{"iopub.status.busy":"2022-04-02T19:54:12.287905Z","iopub.execute_input":"2022-04-02T19:54:12.288476Z","iopub.status.idle":"2022-04-02T19:54:12.301507Z","shell.execute_reply.started":"2022-04-02T19:54:12.288427Z","shell.execute_reply":"2022-04-02T19:54:12.300808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

HELPING FUNCTIONS

","metadata":{}},{"cell_type":"code","source":"# SOD- U2 Net prediction\n################################\ndef 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 > CFG.THRESHOLD] = 1\n predict_np[predict_np <= CFG.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-04-02T19:54:12.303521Z","iopub.execute_input":"2022-04-02T19:54:12.304068Z","iopub.status.idle":"2022-04-02T19:54:12.319025Z","shell.execute_reply.started":"2022-04-02T19:54:12.304032Z","shell.execute_reply":"2022-04-02T19:54:12.318118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

LOADING DATA

","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_csv(\"../input/whale2-cropped-dataset/train2.csv\")\ntest_df = pd.read_csv(\"../input/whale2-cropped-dataset/test2.csv\")\n\ntrain_path = \"../input/whale2-cropped-dataset/cropped_train_images/cropped_train_images\"\ntest_path = \"../input/whale2-cropped-dataset/cropped_test_images/cropped_test_images\"","metadata":{"execution":{"iopub.status.busy":"2022-04-02T19:54:12.320492Z","iopub.execute_input":"2022-04-02T19:54:12.320947Z","iopub.status.idle":"2022-04-02T19:54:12.491375Z","shell.execute_reply.started":"2022-04-02T19:54:12.320867Z","shell.execute_reply":"2022-04-02T19:54:12.490672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\n### Create Kaggle Dataset if not exists \n\nDATASET_NAME = f'happywhale-cropped-removeBackground-v1'\nTRAINING_NAME = f'removedBackground_train_images'\nTESTING_NAME = f'removedBackground_test_image'\n\n!rm -r /tmp/{DATASET_NAME} # remove folder\n\nos.makedirs(f'/tmp/{DATASET_NAME}', exist_ok=True)\nos.makedirs(f'/tmp/{DATASET_NAME}/{TRAINING_NAME}', exist_ok=True)\nos.makedirs(f'/tmp/{DATASET_NAME}/{TESTING_NAME}', exist_ok=True)\n\nwith open('../input/kaggle-json-file/kaggle.json') as f:\n kaggle_creds = json.load(f)\n \nos.environ['KAGGLE_USERNAME'] = kaggle_creds['username']\nos.environ['KAGGLE_KEY'] = kaggle_creds['key']\n\n!kaggle datasets init -p /tmp/{DATASET_NAME}\n\nwith open(f'/tmp/{DATASET_NAME}/dataset-metadata.json') as f:\n dataset_meta = json.load(f)\ndataset_meta['id'] = f'phanttan/{DATASET_NAME}'\ndataset_meta['title'] = DATASET_NAME\nwith open(f'/tmp/{DATASET_NAME}/dataset-metadata.json', \"w\") as outfile:\n json.dump(dataset_meta, outfile)\nprint(dataset_meta)\n\n!cp /tmp/{DATASET_NAME}/dataset-metadata.json /tmp/{DATASET_NAME}/meta.json\n!ls /tmp/{DATASET_NAME}\n\n!kaggle datasets create -u -p /tmp/{DATASET_NAME} ","metadata":{"execution":{"iopub.status.busy":"2022-04-02T19:54:12.492718Z","iopub.execute_input":"2022-04-02T19:54:12.492952Z","iopub.status.idle":"2022-04-02T19:54:17.133533Z","shell.execute_reply.started":"2022-04-02T19:54:12.492921Z","shell.execute_reply":"2022-04-02T19:54:17.132753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_directory = [train_path + '/' + file for file in train_df.image]\ntest_directory = [test_path + '/' + file for file in test_df.image]","metadata":{"execution":{"iopub.status.busy":"2022-04-02T19:54:17.135325Z","iopub.execute_input":"2022-04-02T19:54:17.13567Z","iopub.status.idle":"2022-04-02T19:54:17.20318Z","shell.execute_reply.started":"2022-04-02T19:54:17.135627Z","shell.execute_reply":"2022-04-02T19:54:17.201587Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#

SOD-U2-Net Prediction

","metadata":{}},{"cell_type":"code","source":"%%capture\n\nif CFG.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-04-02T19:54:17.212614Z","iopub.execute_input":"2022-04-02T19:54:17.212821Z","iopub.status.idle":"2022-04-02T19:54:22.679128Z","shell.execute_reply.started":"2022-04-02T19:54:17.212795Z","shell.execute_reply":"2022-04-02T19:54:22.67839Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for idx, img in enumerate(train_directory): \n image, im_orig, sal_map, mask = pred_unet(net, [train_directory[idx]])\n# print(idx)\n try: \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 if (ymin != ymax) & (xmin != xmax):\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 else:\n crop_img = im_orig\n except IndexError:\n crop_img = im_orig\n im = Img.fromarray(crop_img)\n im.save(f'/tmp/{DATASET_NAME}/{TRAINING_NAME}/{train_df.image[idx]}')\n","metadata":{"execution":{"iopub.status.busy":"2022-04-02T19:54:22.68056Z","iopub.execute_input":"2022-04-02T19:54:22.680807Z","iopub.status.idle":"2022-04-02T22:52:47.034505Z","shell.execute_reply.started":"2022-04-02T19:54:22.680772Z","shell.execute_reply":"2022-04-02T22:52:47.033641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for idx, img in enumerate(test_directory): \n image, im_orig, sal_map, mask = pred_unet(net, [test_directory[idx]])\n try: \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 if (ymin != ymax) & (xmin != xmax):\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 else:\n crop_img = im_orig\n except IndexError:\n crop_img = im_orig\n im = Img.fromarray(crop_img)\n im.save(f'/tmp/{DATASET_NAME}/{TESTING_NAME}/{test_df.image[idx]}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"version_name = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n!kaggle datasets version -m {version_name} -p /tmp/{DATASET_NAME} -r zip -q","metadata":{"execution":{"iopub.status.busy":"2022-04-03T00:32:54.654702Z","iopub.execute_input":"2022-04-03T00:32:54.654917Z","iopub.status.idle":"2022-04-03T00:34:05.05519Z","shell.execute_reply.started":"2022-04-03T00:32:54.65489Z","shell.execute_reply":"2022-04-03T00:34:05.054308Z"},"trusted":true},"execution_count":null,"outputs":[]}]}