{"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":"## Enhance under water images with U-shape transformer without resizing image! 👀\n\n### [U-shape Transformer Underwater Image Enhancement]\n\npaper: https://arxiv.org/abs/2111.11843\n\ncode (pytorch): https://github.com/LintaoPeng/U-shape_Transformer\n\nBefore I start with U-shape Transformer, I found out that this code resize the output image into 256x256 which is not efficient when training with detection code.\n\nSo I edit some code lines to handle images by croping into 256 x 256 patches and preprocess each piece.  \n\nI refer to code provided in korean competition site \"DACON\" which I linked below!\n\nreference: https://dacon.io/competitions/official/235746/codeshare/2874?page=2&dtype=recent","metadata":{}},{"cell_type":"code","source":"\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%cd /kaggle/working/\n!git clone https://github.com/LintaoPeng/U-shape_Transformer_for_Underwater_Image_Enhancement.git\n%cd /kaggle/working/U-shape_Transformer_for_Underwater_Image_Enhancement","metadata":{"execution":{"iopub.status.busy":"2022-01-15T11:45:43.122385Z","iopub.execute_input":"2022-01-15T11:45:43.122662Z","iopub.status.idle":"2022-01-15T11:45:43.128635Z","shell.execute_reply.started":"2022-01-15T11:45:43.122627Z","shell.execute_reply":"2022-01-15T11:45:43.127951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -r requirements.txt","metadata":{"execution":{"iopub.status.busy":"2022-01-15T11:47:29.623161Z","iopub.execute_input":"2022-01-15T11:47:29.623866Z","iopub.status.idle":"2022-01-15T11:47:38.050888Z","shell.execute_reply.started":"2022-01-15T11:47:29.623828Z","shell.execute_reply":"2022-01-15T11:47:38.050042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### download saved_model.zip folder where models saved and place it where you want\n\nhttps://drive.google.com/file/d/19a_kDJTT5S96kzwQntEMhSxAPYw4xY2P/view","metadata":{}},{"cell_type":"code","source":"import os\n#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\nimport argparse\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.utils as utils\nimport time\nfrom torch.autograd import Variable\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\nfrom glob import glob\nfrom torch.nn.modules.loss import _Loss\nfrom net.Ushape_Trans import *\n#from dataset import prepare_data, Dataset\nfrom net.utils import *\nimport cv2\nimport matplotlib.pyplot as plt\nfrom utility import plots as plots, ptcolor as ptcolor, ptutils as ptutils, data as data\nfrom loss.LAB import *\nfrom loss.LCH import *\nfrom torchvision.utils import save_image\n\nos.environ[\"KMP_DUPLICATE_LIB_OK\"]=\"TRUE\"","metadata":{"execution":{"iopub.status.busy":"2022-01-15T11:47:38.052744Z","iopub.execute_input":"2022-01-15T11:47:38.053024Z","iopub.status.idle":"2022-01-15T11:47:38.082595Z","shell.execute_reply.started":"2022-01-15T11:47:38.052986Z","shell.execute_reply":"2022-01-15T11:47:38.081797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dtype = 'float32'\n\ntorch.set_default_tensor_type(torch.FloatTensor)\n\n\ndef split(img):\n    output=[]\n    output.append(F.interpolate(img, scale_factor=0.125))\n    output.append(F.interpolate(img, scale_factor=0.25))\n    output.append(F.interpolate(img, scale_factor=0.5))\n    output.append(img)\n    return output\n\n# Initialize generator\ngenerator = Generator().cuda()\ngenerator.load_state_dict(torch.load(\"/kaggle/input/models/saved_models/G/generator_795.pth\"))\n\n\nimg_size = 256\npath = '/kaggle/input/tensorflow-great-barrier-reef/train_images/video_0' # data directory\n\n\n\ndef predict(img_paths, stride=32, batch_size=1):\n    results = []\n    for img_path in os.listdir(img_paths):\n        img = cv2.imread(os.path.join(img_paths, img_path))\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        #img = np.array(img).astype(dtype)\n        #img = torch.from_numpy(img)  # tensor (h, w, c)\n\n\n        crop = []\n        position = []\n        batch_count = 0\n\n        result_img = np.zeros_like(img)\n        voting_mask = np.zeros_like(img)\n        i = 0\n        for top in tqdm(range(0, img.shape[0], stride)):\n            for left in range(0, img.shape[1], stride):\n                piece = np.zeros([img_size, img_size, 3], np.float32)\n                temp = img[top:top + img_size, left:left + img_size, :]\n                piece[:temp.shape[0], :temp.shape[1], :] = temp\n                crop.append(piece)\n                position.append([top, left])\n                batch_count += 1\n                if batch_count == batch_size:\n\n\n                    crop = np.array(crop).astype(dtype)\n                    crop = torch.from_numpy(crop)  #tensor (b, h, w, c)\n                    crop = crop.permute(0, 3, 1, 2) # tensor (b, c, h, w)\n                    #.unsqueeze(0) => make batch 1\n                    crop = crop / 255.0\n                    crop = Variable(crop).cuda()\n                    output = generator(crop)*255\n                    pred = output[3].data\n\n                    # check cropped image output\n                    save_image(pred, \"./test/output/cropped_\" + img_path, nrow=5, normalize=True)\n\n                    crop = []\n                    batch_count = 0\n                    for num, (t, l) in enumerate(position):\n                        piece = pred[num]\n                        piece = piece.permute(1, 2, 0)\n                        h, w, c = result_img[t:t + img_size, l:l + img_size, :].shape\n                        #result_img = torch.Tensor(result_img).cuda()\n                        piece = piece.cpu().detach().numpy()\n                        result_img = result_img.astype(dtype)\n                        result_img[t:t + img_size, l:l + img_size, :] += piece[:h, :w, :]\n                        voting_mask[t:t + img_size, l:l + img_size, :] += 1\n                    position = []\n\n\n        result_img = result_img / voting_mask\n        #result_img = result_img.astype(np.uint8)\n        #results.append(result_img)\n        results = np.array(result_img).astype(dtype)\n        results = torch.from_numpy(results)\n        results = results.permute(2, 0, 1).unsqueeze(0)# tensor (b, c, h, w)\n        # save image to your output folder dir\n        save_image(results, \"../outputdir/\" + img_path, nrow=5, normalize=True)\n\n    return results\n\n\nif __name__ == \"__main__\":\n\n    predict(path)","metadata":{"execution":{"iopub.status.busy":"2022-01-15T11:47:42.675819Z","iopub.execute_input":"2022-01-15T11:47:42.676392Z","iopub.status.idle":"2022-01-15T11:47:49.048755Z","shell.execute_reply.started":"2022-01-15T11:47:42.676353Z","shell.execute_reply":"2022-01-15T11:47:49.047757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}