{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Verifiying results from https://arxiv.org/pdf/1911.05140v1.pdf"},{"metadata":{},"cell_type":"markdown","source":"# Unsupervised medical images segmentation using generative adversarial networks: from edge diagrams to segmentation maps"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pathlib\nimport os\nimport sys\nimport time\nimport h5py\nimport random\n\nimport numpy as np\nimport tensorflow as tf\nimport keras\nfrom keras import losses\nfrom keras import models\nimport matplotlib.pyplot as plt\n\nfrom skimage import color, exposure, io, img_as_float, transform, filters, morphology, measure\nfrom sklearn import metrics\nfrom PIL import Image\n\nimport cv2\nfrom shutil import copyfile, make_archive, copy\n\ntraining_input_dir = os.path.join(\"../input/isic2018train\", \"ISIC2018_Task1-2_Training_Input\")\ntraining_ground_truth_dir = os.path.join(\"D:/\", \"Dropbox\", \"Datasets\", \"ISIC2018\", \"ISIC2018_Task1_Training_GroundTruth\")\n# richer convolutional features with non-maximum supression\ntraining_rcf_nms_dir = os.path.join(\"../input/richer-convolutional-features\", \"isic_steps\", \"train_rcf_nms\")\nvalidation_input_dir = os.path.join(\"D:/\", \"Dropbox\", \"Datasets\", \"ISIC2018\", \"ISIC2018_Task1-2_Validation_Input\")\nvalidation_rcf_nms_dir = os.path.join(\"../input/richer-convolutional-features\", \"isic_steps\", \"validate_rcf_nms\")\ntest_input_dir = os.path.join(\"../input/siim-isic-melanoma-classification\", \"jpeg\", \"test\")\n\n# %%\ndef load_images(filepath):    \n    '''\n    Loads the available images.\n    '''\n    data_root = pathlib.Path(os.path.join(filepath))\n    image_paths = [str(path) for path in  list(data_root.glob('*'))]\n    return image_paths\n\n\ndef preprocess_images(image_paths, savedir, load_prev=False):\n    '''\n    Preprocessing for ISIC images\n    '''\n    if not load_prev or not os.path.isdir(savedir):\n        os.makedirs(savedir, exist_ok=True)\n        for path in image_paths[:300]:\n            savepath = os.path.join(savedir, os.path.basename(path))\n            print(savepath)\n            im = img_as_float(io.imread(path))        \n            im = transform.resize(im, (256, 256), mode=\"reflect\", anti_aliasing=True)\n            im = np.uint8(im * 255)\n            io.imsave(fname=savepath,arr=im)\n        \n    all_image_paths = [str(path) for path in  list(pathlib.Path(savedir).glob('*'))]\n    return all_image_paths\n\n\ndef shrink_bool_img(image, new_size):\n    im = np.zeros((new_size, new_size), dtype=np.bool)\n    big = image.shape[0]\n    small = im.shape[0]\n    scale = big // small\n    for i in np.arange(1, small-1):\n        for j in np.arange(1, small-1):\n            im[i, j] = np.any(\n                image[i*scale:i*scale+scale, j*scale:j*scale+scale])\n    return im\n\n\ndef grow_bool_img(image, new_size):\n    im = np.empty((new_size, new_size), dtype=np.bool)\n    big = im.shape[0]\n    small = image.shape[0]\n    scale = big // small\n    for i in np.arange(small):\n        for j in np.arange(small):\n            im[i*scale:i*scale+scale, j*scale:j*scale +\n                scale] = np.full((scale, scale), image[i, j])\n    return im\n\n\ndef convert_to_edge_diagrams(image_paths_rcf, savedir, load_prev=True):\n    '''\n    Converts the RCF + NMS output into edge diagrams.\n    '''    \n    if not load_prev or not os.path.isdir(savedir):\n        os.makedirs(savedir, exist_ok=True)\n        for path_rcf in image_paths_rcf:\n            savepath = os.path.join(savedir, os.path.basename(path_rcf))\n            \n            img_rcf = color.rgb2gray(img_as_float(io.imread(path_rcf)))\n        \n            # Remove the outline of the cone\n            img_rcf = cv2.resize(img_rcf, (272, 272))\n            img_rcf = img_rcf[8:264, 8:264]\n            img_rcf = img_rcf >= filters.threshold_otsu(img_rcf)\n            \n            # Shrink and simplify\n            img_rcf = shrink_bool_img(img_rcf, 32)\n            label_rcf = measure.label(img_rcf)\n            for region in measure.regionprops(label_rcf):\n                if region.area < 3:\n                    for x, y in region.coords:\n                        img_rcf[x, y] = False\n            img_rcf = morphology.skeletonize(img_rcf)\n            \n            # save the output edge diagram\n            edge_diagram = (img_rcf*255).astype(np.uint8)            \n            io.imsave(fname=savepath, arr=edge_diagram)\n                \n    out_image_paths = [str(path) for path in  list(pathlib.Path(savedir).glob('*'))]\n    return out_image_paths\n\n\ndef create_synthetic_edge_diagrams(savedir_mask, savedir_edge, n_synth=300, load_prev=True):\n    '''\n    Generates the 'ground truth' segmentation masks for generating synthetic\n    ultrasound images, and combines them with the generated cones.\n    '''\n    if not load_prev or not os.path.isdir(savedir_mask) or not os.path.isdir(savedir_edge):\n        os.makedirs(savedir_mask, exist_ok=True)\n        os.makedirs(savedir_edge, exist_ok=True)\n        \n    for n in np.arange(n_synth):\n        print('Creating {} synthetic edge diagrams'.format(n_synth))\n        savepath_mask = os.path.join(savedir_mask, '{}.png'.format(n))\n        savepath_edge = os.path.join(savedir_edge, '{}.png'.format(n))\n                     \n        # Draw a big ellipse (the lesion outline) with random properties\n        lesion_outline = np.zeros((32, 32))\n        minor_axis = np.random.randint(2, 10)\n        major_axis = np.random.randint(minor_axis + 1, minor_axis + 8)\n        center_x = np.random.randint(12, 20)\n        center_y = np.random.randint(12, 20)\n        rotation = np.random.randint(-60, 60)\n    \n        # Add some random dots within the outline\n        lesion = np.zeros_like(lesion_outline, dtype=np.uint8)\n        cv2.ellipse(lesion, (center_x, center_y),\n                    (minor_axis, major_axis), rotation, 0, 360, 1, -1)\n        x, y = np.where(lesion)\n        idx = np.random.randint(len(x), size=np.random.randint(minor_axis))\n        lesion_outline[x[idx], y[idx]] = True\n        \n        # draw half of the ellipses in arcs, and half fully, just for some variation\n        random_number = np.random.randint(2)\n        if random_number == 0:\n            # draw circle in arcs\n            angles, step = np.linspace(0, 360, np.random.randint(5, 8), retstep=True)\n            for angle in angles:\n                cv2.ellipse(lesion_outline,\n                    (center_x, center_y),\n                    (minor_axis, major_axis),\n                    rotation,\n                    angle,\n                    angle + step // 2,\n                    1,\n                    1,\n                )\n        else:\n            # draw circle fully, and add some decorations\n            cv2.ellipse(\n                lesion_outline,\n                (center_x, center_y),\n                (minor_axis, major_axis),\n                rotation,\n                0,\n                360,\n                1,\n                1,\n            )\n            # if the lesion isn't too big, maybe draw some hairs\n            if (minor_axis <= 6 and major_axis <= 8):\n                random_number = np.random.randint(4)\n                if random_number == 0:\n                    n_hairs = np.random.randint(3, 6)\n                    for _ in np.arange(n_hairs):\n                        start_x = np.random.randint(2, 30)\n                        start_y = np.random.randint(2, 30)\n                        stop_x = start_x + np.random.randint(6, 10) * (1 if np.random.random() < 0.5 else -1)\n                        stop_y = start_y + np.random.randint(6, 10) * (1 if np.random.random() < 0.5 else -1)\n                        cv2.line(\n                            lesion_outline,\n                            (start_x, start_y),\n                            (stop_x, stop_y),\n                            1,\n                        )\n            \n                elif (minor_axis <= 4 and major_axis <= 6):\n                    # if the lesion is small, maybe draw some other things that appear in the training set\n                    random_number = np.random.randint(8)\n                    if random_number == 0:\n                        # draw the 2 circles on either side\n                        for circle_x in [0, 32]:\n                            cv2.ellipse(\n                                lesion_outline,\n                                (circle_x, 16),\n                                (4, 10),\n                                0,\n                                0,\n                                360,\n                                1,\n                                1,\n                            )\n                    elif random_number == 1:\n                        # draw the 4 pen blots around the lesion\n                        for delta_x in [8, -8]:\n                            for delta_y in [8, -8]:\n                                cv2.ellipse(\n                                    lesion_outline,\n                                    (center_x - delta_x, center_y - delta_y),\n                                    (1, 1),\n                                    0,\n                                    0,\n                                    360,\n                                    1,\n                                    1,\n                                )\n                    elif random_number == 2:\n                        # draw a small ruler\n                        cv2.line(\n                            lesion_outline,\n                            (np.random.randint(8, 12), np.random.randint(28, 30)),\n                            (np.random.randint(20, 24), np.random.randint(28, 30)),\n                            1,\n                        )\n                    elif random_number == 3:\n                        # draw a big ruler\n                        cv2.line(\n                            lesion_outline,\n                            (0, np.random.randint(26, 32)),\n                            (32, np.random.randint(26, 32)),\n                            1,\n                        )\n                    elif random_number == 4:\n                        # draw the dermoscope lens\n                        cv2.ellipse(\n                            lesion_outline,\n                            (16, 16),\n                            (np.random.randint(18, 20), np.random.randint(18, 20)),\n                            0,\n                            0,\n                            360,\n                            1,\n                            1,\n                        )\n\n        mask = grow_bool_img(lesion, 256)\n        edge_diagram = grow_bool_img(lesion_outline, 256)\n        \n        mask = mask.astype(np.uint8) * 255\n        edge_diagram = edge_diagram.astype(np.uint8)\n        \n        io.imsave(fname=savepath_mask, arr=mask)\n        io.imsave(fname=savepath_edge, arr=edge_diagram)\n        \n    image_paths_mask = [str(path) for path in list(pathlib.Path(savedir_mask).glob('*'))]\n    image_paths_edge = [str(path) for path in  list(pathlib.Path(savedir_edge).glob('*'))]\n    return image_paths_mask, image_paths_edge\n\n\ndef use_mask_rcnn(train_path, mask_path, test_path, test_mask, savedir):\n    '''\n    Trains the mask_rcnn model on given images.\n    '''\n    # Train MRCNN\n    mrcnn = MRCNN(train_path, mask_path, val_path, savedir='Models/MaskRCNN')\n    mrcnn.create_model()\n    mrcnn.train()\n        \n    return preds, scores","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# %%\n# Load all of the available images\nprint('Loading all images... ')\nstart = time.time()\n# train_image_paths = load_images(training_input_dir)\n# test_image_paths = load_images(test_input_dir)\nend = time.time()\nprint('took {} seconds'.format(end-start))\n\n# # Step 1: Richer convolutional featuers (RCF) and Non-Maximum Suppression\n# # This step requires pretrained models in caffe and matlab, so must be done externally.\n# # See paper for details, as we just use two pre-packaged libraries unaltered.\n# # For this dataset, we don't preprocess, so we just resize the images in the RCF+NMS code\ntrain_rcfnms_paths = [str(path) for path in list(pathlib.Path(training_rcf_nms_dir).glob('*'))]\nval_rcfnms_paths = [str(path) for path in list(pathlib.Path(validation_rcf_nms_dir).glob('*'))]\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # Step 2: Convert RCF + NMS output to edge diagrams\nprint('Convert RCF+NMS output to edge diagrams... ')\nstart = time.time()\ntrain_edge_paths = convert_to_edge_diagrams(train_rcfnms_paths, savedir='train_edge', load_prev=True)\nval_edge_paths = convert_to_edge_diagrams(val_rcfnms_paths, savedir='val_edge', load_prev=True)\nend = time.time()\nprint('took {} seconds'.format(end-start))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm synth_mask/* \n!rm synth_edge/*","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # Step 3: Generate synthetic edge diagrams\nprint('Creating synthetic edge diagrams... ')\nstart = time.time()\nsynth_mask_paths, synth_edge_paths = create_synthetic_edge_diagrams('synth_mask', 'synth_edge', n_synth=300, load_prev=True)\nend = time.time()\nprint('took {} seconds'.format(end-start))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# !rm -rf train_img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# !cp -r ../input/isic2018train/ISIC2018_Task1-2_Training_Input ./train_img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install dominate","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"preprocess_images(load_images(training_input_dir), 'train_img', load_prev=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm train_img/*","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# !mv train_edge train_label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # Step 4: Train pix2pixHD and choose the best epoch using the Frechet Inception Distance\n# # This uses pytorch instead of tensorflow, so we run these steps externally using the pix2pixHD code provided by nVidia\nprint('Training pix2pixhD with only training images... ')\nstart = time.time()\n# # Train command: python train.py --name pix2ultra --resize_or_crop none --checkpoints_dir pix2ultra/checkpoints --dataroot pix2ultra/datasets/isic2018/ --nThreads 4 --display_winsize 256 --tf_log --no_instance --label_nc 2\n!python ../input/pix2pix/pix2pixHD-master/train.py --name pix2ultra --resize_or_crop none --checkpoints_dir pix2ultra/checkpoints --dataroot ./ --nThreads 1 --display_winsize 256 --tf_log --no_instance --label_nc 2\n# # FID command: python fid.py path/to/images path/to/other/images --gpu 0\nend = time.time()\nprint('took {} seconds'.format(end-start))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # Step 5: Use the trained pix2pixHD model to generate synthetic ultrasound images from the synthetic edge diagrams\n# print(' all images... ')\n# start = time.time()\n# # Test command: python test.py --name pix2ultra --resize_or_crop none --checkpoints_dir pix2ultra/checkpoints --results_dir pix2ultra/results --how_many 3000 --dataroot pix2ultra/datasets/isic2018/ --display_winsize 256 --no_instance --label_nc 2\n# end = time.time() \n# print('took {} seconds'.format(end-start))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # Step 6: Use Mask-RCNN to train on the synthetic images - test is done on competiion submission page\n# print(' all images... ')\n# start = time.time()\n# train_path = 'synth_isic'\n# mask_path = 'synth_mask'\n# unet_preds, unet_metrics = use_mask_rcnn(train_path, mask_path, savedir='mrcnn_masks')\n# end = time.time()\n# print('took {} seconds'.format(end-start))","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}