{"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":"code","source":"# https://www.tensorflow.org/tfmodels/vision/instance_segmentation","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -U -q \"tf-models-official\" > /dev/null\n!pip install -U -q remotezip tqdm opencv-python einops > /dev/null","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport cv2\nimport io\nimport glob\nimport json\nimport shutil\nimport pprint\nimport pathlib\nimport tempfile\nimport requests\nimport collections\nimport matplotlib\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport tifffile as tiff\nimport matplotlib.pyplot as plt\n\nfrom PIL import Image\nfrom six import BytesIO\nfrom etils import epath\nfrom IPython import display\nfrom urllib.request import urlopen\nfrom pycocotools import mask as coco_mask\nfrom scipy.ndimage import label\nfrom tqdm import tqdm\nfrom sklearn.model_selection import KFold\nfrom dataclasses import dataclass\nfrom typing import List, Union\n\nimport orbit\nimport tensorflow as tf\nimport tensorflow_models as tfm\nimport tensorflow_datasets as tfds\n\nfrom official.core import exp_factory\nfrom official.core import config_definitions as cfg\nfrom official.vision.data import tfrecord_lib\nfrom official.vision.serving import export_saved_model_lib\nfrom official.vision.dataloaders.tf_example_decoder import TfExampleDecoder\nfrom official.vision.utils.object_detection import visualization_utils\nfrom official.vision.ops.preprocess_ops import normalize_image, resize_and_crop_image\nfrom official.vision.data.create_coco_tf_record import coco_annotations_to_lists\n\npp = pprint.PrettyPrinter(indent=4) # Set Pretty Print Indentation\nprint(tf.__version__) # Check the version of tensorflow used\n\n%matplotlib inline","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@dataclass\nclass BaseConfig:\n    seed: int\n    num_fold: int\n    group_col: str\n    target_col: Union[str, List[str]]\n        \nclass CFG(BaseConfig):\n    seed = 42\n    num_fold = 5\n    target_col = 'id'\n    \n\ndef split_kfold(df: pd.DataFrame, cfg: BaseConfig) -> pd.DataFrame:\n    df[\"kfold\"] = -1\n    kf = KFold(n_splits=cfg.num_fold, shuffle=True, random_state=cfg.seed)\n    for n, (trn_index, val_index) in enumerate(kf.split(X=df, y=df[cfg.target_col])):\n        df.loc[val_index, \"kfold\"] = int(n)\n    df[\"kfold\"] = df[\"kfold\"].astype(int)\n    return df\n\n\ndef generate_chunks(data_list, num_chunk):\n    for i in range(0, len(data_list), num_chunk):\n        yield data_list[i:i + num_chunk]\n        \n\ndef _convert_to_feature(\n    feature0, feature1, feature2, \n    feature3, feature4, feature5,\n    feature6, feature7, feature8,\n    feature9, feature10, feature11,\n):\n    feature = {\n        'image/source_id': tfrecord_lib.convert_to_feature(feature0),\n        'image/encoded': tfrecord_lib.convert_to_feature(feature1),\n        'image/height': tfrecord_lib.convert_to_feature(feature2),\n        'image/width': tfrecord_lib.convert_to_feature(feature3),\n        'image/object/class/label': tfrecord_lib.convert_to_feature(feature4),\n        'image/object/is_crowd': tfrecord_lib.convert_to_feature(feature5),\n        'image/object/area': tfrecord_lib.convert_to_feature(feature5, 'float_list'),\n        'image/object/bbox/xmin': tfrecord_lib.convert_to_feature(feature7, 'float_list'),\n        'image/object/bbox/ymin': tfrecord_lib.convert_to_feature(feature8, 'float_list'),\n        'image/object/bbox/xmax': tfrecord_lib.convert_to_feature(feature9, 'float_list'),\n        'image/object/bbox/ymax': tfrecord_lib.convert_to_feature(feature10, 'float_list'),\n        'image/object/mask': tfrecord_lib.convert_to_feature(feature11),\n    }\n    return feature\n\n\ndef get_coco_bounding_box_from_mask(mask):\n    image_width = 512\n    image_height = 512\n\n    height, width = mask.shape[:2]\n\n    non_zero_indices = np.nonzero(mask)\n\n    if len(non_zero_indices[0]) == 0:\n        return None\n\n    y_min = np.min(non_zero_indices[0])\n    y_max = np.max(non_zero_indices[0])\n    x_min = np.min(non_zero_indices[1])\n    x_max = np.max(non_zero_indices[1])\n\n    x_min /= image_width\n    y_min /= image_height\n    x_max /= image_width\n    y_max /= image_height\n    \n    coco_bbox = [\n        float(x_min),\n        float(y_min),\n        float(x_max),\n        float(y_max)\n    ]\n    return coco_bbox\n\n\ndef convert_polygon_to_segmentation(polygon):\n    segmentation = []\n    for x, y in polygon:\n        segmentation.append(float(x))\n        segmentation.append(float(y))\n    return [segmentation]\n\n\ndef coco_segmentation_to_mask_png(segmentation, height, width, is_crowd):\n    \"\"\"Encode a COCO mask segmentation as PNG string.\"\"\"\n    run_len_encoding = coco_mask.frPyObjects(segmentation, height, width)\n    binary_mask = coco_mask.decode(run_len_encoding)\n    if not is_crowd:\n        binary_mask = np.amax(binary_mask, axis=2)\n\n    return tfrecord_lib.encode_mask_as_png(binary_mask)\n\n\ndef coco_segmentation_coordinate_polygon_to_binary_segmentation_mask(cords) -> np.ndarray:\n    # thanks to https://www.kaggle.com/code/itsuki9180/hubmap-making-dataset\n    instance_mask = np.zeros((512, 512), dtype=np.float32)\n    cords = annot['coordinates']\n    for cd in cords:\n        rr, cc = np.array([i[1] for i in cd]), np.asarray([i[0] for i in cd])\n        instance_mask[rr, cc] = 1\n\n    contours,_ = cv2.findContours((instance_mask*255).astype(np.uint8), 1, 2)\n    zero_img = np.zeros([instance_mask.shape[0], instance_mask.shape[1], 1], dtype=\"uint8\")\n    for p in contours:\n        cv2.fillPoly(zero_img, [p], (1, 1, 1))\n    contours, hierarchy = cv2.findContours(instance_mask.astype(\"uint8\"), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)\n    img_with_area = zero_img\n    for i in range(len(contours)):\n        cv2.fillPoly(img_with_area, [contours[i][:,0,:]], (1,1,1), lineType=cv2.LINE_8, shift=0)\n    return img_with_area","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open('/kaggle/input/hubmap-hacking-the-human-vasculature/polygons.jsonl', 'r') as json_file:\n    json_list = list(json_file)\n    \ntile_meta_df = pd.read_csv(\"/kaggle/input/hubmap-hacking-the-human-vasculature/tile_meta.csv\")\ntile_meta_df = split_kfold(tile_meta_df, CFG)\n\ndf = pd.DataFrame(json_list, columns=['annot_dict'])\ndf['id'] = df['annot_dict'].map(lambda x: eval(x)['id'])\ndf = pd.merge(df, tile_meta_df, on='id', how='inner')\nprint(df.shape)\ndf.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create tfrecord","metadata":{}},{"cell_type":"code","source":"tf_dataset_dir = '/tmp/hubmap-5fold-k-split'\n\nif not os.path.exists(tf_dataset_dir):\n    os.mkdir(tf_dataset_dir)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"multi_idx = df[['source_wsi', 'dataset', 'kfold']].value_counts().keys()\n\nid_all = 0\nfor source_wsi, dataset, kfold in multi_idx:\n    df_subset = df.query(f\"kfold == {kfold} & dataset == {dataset} & source_wsi == {source_wsi}\").reset_index(drop=True)\n    \n    file_path = f'/tmp/hubmap-5fold-k-split/fold{kfold}-dataset{dataset}-sourcewsi{source_wsi}-{len(df_subset)}.tfrecords'\n    with tf.io.TFRecordWriter(\n        file_path,\n    ) as writer:\n        for idx, annot_dict in tqdm(zip(\n                    df_subset['id'].values,\n                    df_subset['annot_dict'].values,\n                ), total=len(df_subset)):\n\n            tldc = eval(annot_dict)\n            \n            source_id = str(id_all).encode('utf8')\n\n            array = tiff.imread(f'/kaggle/input/hubmap-hacking-the-human-vasculature/train/{idx}.tif')\n            img_example = Image.fromarray(array).convert(\"RGB\")\n\n            img = np.array(img_example)\n            img = tf.io.encode_jpeg(img).numpy()\n\n            height = 512\n            width = 512\n            \n            groundtruth_classes = []\n            groundtruth_is_crowd = []\n            groundtruth_area = []\n            groundtruth_xmin = []\n            groundtruth_ymin = []\n            groundtruth_xmax = []\n            groundtruth_ymax = []\n            groundtruth_instance_masks_png = []\n            for annot in tldc['annotations']:\n                if annot['type'] == 'blood_vessel':\n                    poly = annot['coordinates']\n                    binary_segmentation_mask = coco_segmentation_coordinate_polygon_to_binary_segmentation_mask(poly)\n                    groundtruth_box = get_coco_bounding_box_from_mask(binary_segmentation_mask)\n                    mask_area = float(binary_segmentation_mask.sum().sum())\n\n                    segm = convert_polygon_to_segmentation(poly[0])\n                    mask_png = coco_segmentation_to_mask_png(segm, 512, 512, 0)\n                    groundtruth_instance_masks_png.append(mask_png)\n                    groundtruth_classes.append(1)\n                    groundtruth_is_crowd.append(False)\n                    groundtruth_area.append(mask_area)\n                    groundtruth_xmin.append(groundtruth_box[0])\n                    groundtruth_ymin.append(groundtruth_box[1])\n                    groundtruth_xmax.append(groundtruth_box[2])\n                    groundtruth_ymax.append(groundtruth_box[3])  \n            \n            if len(groundtruth_classes) == 0:\n                continue\n            feature = _convert_to_feature(\n                feature0 = source_id,\n                feature1 = img,\n                feature2 = height,\n                feature3 = width,\n                feature4 = groundtruth_classes,\n                feature5 = groundtruth_is_crowd,\n                feature6 = groundtruth_area,\n                feature7 = groundtruth_xmin,\n                feature8 = groundtruth_ymin,\n                feature9 = groundtruth_xmax,\n                feature10 = groundtruth_ymax,\n                feature11 = groundtruth_instance_masks_png,\n            )\n\n            id_all += 1\n\n            example = tf.train.Example(features=tf.train.Features(feature=feature))\n            serialized_example = example.SerializeToString()\n            writer.write(serialized_example) ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf_ex_decoder = TfExampleDecoder(include_mask=True)\n\nds = tf.data.TFRecordDataset(\n    glob.glob('/tmp/hubmap-5fold-k-split/*.tfrecords'),\n    num_parallel_reads=tf.data.AUTOTUNE, \n)\n\n\nfor i, serialized_example in enumerate(ds):\n    decoded_tensors = tf_ex_decoder.decode(serialized_example)\n    if len(decoded_tensors['groundtruth_boxes'].numpy()) > 1:\n        break\n    \nprint(decoded_tensors.keys())\nprint(decoded_tensors['source_id'].numpy())\nprint(decoded_tensors['height'].numpy())\nprint(decoded_tensors['width'].numpy())\nprint(decoded_tensors['groundtruth_classes'].numpy())\nprint(decoded_tensors['groundtruth_is_crowd'].numpy())\nprint(decoded_tensors['groundtruth_area'].numpy())\nprint(decoded_tensors['groundtruth_boxes'].numpy())\nprint(decoded_tensors['groundtruth_instance_masks'].numpy().shape)\nprint(decoded_tensors['groundtruth_instance_masks_png'].numpy().shape)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(decoded_tensors['image'].numpy())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"msk = 0\nfor m in decoded_tensors['groundtruth_instance_masks'].numpy():\n    msk += m\nplt.imshow(msk)\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"fold = 0\n\n# only dataset 1\nvalid_data_input_path = glob.glob(f'/tmp/hubmap-5fold-k-split/fold{fold}-dataset1-*.tfrecords') # VALID_FILENAMES\ntrain_data_input_path = glob.glob('/tmp/hubmap-5fold-k-split/fold*-dataset1-*.tfrecords') # TRAIN_FILENAMES\n\n\ntrain_data_input_path = list(set(train_data_input_path) - set(valid_data_input_path))\nprint(train_data_input_path, valid_data_input_path)\n\nmodel_dir = './trained_model/'\nexport_dir ='./exported_model/'\n\nif not os.path.exists(model_dir):\n    os.mkdir(model_dir)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"exp_config = exp_factory.get_exp_config('maskrcnn_mobilenet_coco')\n\n\nmodel_ckpt_path = './model_ckpt/'\nif not os.path.exists(model_ckpt_path):\n    os.mkdir(model_ckpt_path)\n\n!gsutil cp gs://tf_model_garden/vision/mobilenet/v2_1.0_float/ckpt-180648.data-00000-of-00001 './model_ckpt/'\n!gsutil cp gs://tf_model_garden/vision/mobilenet/v2_1.0_float/ckpt-180648.index './model_ckpt/'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NUM_CLASSES = 1\nBATCH_SIZE = 4\nHEIGHT, WIDTH = 512, 512\nIMG_SHAPE = [HEIGHT, WIDTH, 3]\n\n\n# Backbone Config\nexp_config.task.annotation_file = None\nexp_config.task.freeze_backbone = True\nexp_config.task.init_checkpoint = \"./model_ckpt/ckpt-180648\"\nexp_config.task.init_checkpoint_modules = \"backbone\"\n\n# Model Config\nexp_config.task.model.num_classes = NUM_CLASSES + 1\nexp_config.task.model.input_size = IMG_SHAPE\n\n# Training Data Config\nexp_config.task.train_data.input_path = train_data_input_path\nexp_config.task.train_data.dtype = 'float32'\nexp_config.task.train_data.global_batch_size = BATCH_SIZE\nexp_config.task.train_data.shuffle_buffer_size = 64\nexp_config.task.train_data.parser.aug_scale_max = 1.0\nexp_config.task.train_data.parser.aug_scale_min = 1.0\n\n# Validation Data Config\nexp_config.task.validation_data.input_path = valid_data_input_path\nexp_config.task.validation_data.dtype = 'float32'\nexp_config.task.validation_data.global_batch_size = BATCH_SIZE\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"logical_device_names = [logical_device.name for logical_device in tf.config.list_logical_devices()]\n\nif 'GPU' in ''.join(logical_device_names):\n  print('This may be broken in Colab.')\n  device = 'GPU'\nelif 'TPU' in ''.join(logical_device_names):\n  print('This may be broken in Colab.')\n  device = 'TPU'\nelse:\n  print('Running on CPU is slow, so only train for a few steps.')\n  device = 'CPU'\n\n\ntrain_steps = 2000\nexp_config.trainer.steps_per_loop = 200 # steps_per_loop = num_of_training_examples // train_batch_size\n\nexp_config.trainer.summary_interval = 200\nexp_config.trainer.checkpoint_interval = 200\nexp_config.trainer.validation_interval = 200\nexp_config.trainer.validation_steps =  200 # validation_steps = num_of_validation_examples // eval_batch_size\nexp_config.trainer.train_steps = train_steps\nexp_config.trainer.optimizer_config.warmup.linear.warmup_steps = 200\nexp_config.trainer.optimizer_config.learning_rate.type = 'cosine'\nexp_config.trainer.optimizer_config.learning_rate.cosine.decay_steps = train_steps\nexp_config.trainer.optimizer_config.learning_rate.cosine.initial_learning_rate = 0.07\nexp_config.trainer.optimizer_config.warmup.linear.warmup_learning_rate = 0.05","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Setting up the Strategy\nif exp_config.runtime.mixed_precision_dtype == tf.float16:\n    tf.keras.mixed_precision.set_global_policy('mixed_float16')\n\nif 'GPU' in ''.join(logical_device_names):\n  distribution_strategy = tf.distribute.MirroredStrategy()\nelif 'TPU' in ''.join(logical_device_names):\n  tf.tpu.experimental.initialize_tpu_system()\n  tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='/device:TPU_SYSTEM:0')\n  distribution_strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n  print('Warning: this will be really slow.')\n  distribution_strategy = tf.distribute.OneDeviceStrategy(logical_device_names[0])\n\nprint(\"Done\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### visualize sample","metadata":{}},{"cell_type":"code","source":"def show_batch(raw_records, num_of_examples):\n    plt.figure(figsize=(20, 20))\n    use_normalized_coordinates=True\n    min_score_thresh = 0.30\n    for i, serialized_example in enumerate(raw_records):\n        plt.subplot(1, 3, i + 1)\n        decoded_tensors = tf_ex_decoder.decode(serialized_example)\n        image = decoded_tensors['image'].numpy().astype('uint8')\n        scores = np.ones(shape=(len(decoded_tensors['groundtruth_boxes'])))\n        visualization_utils.visualize_boxes_and_labels_on_image_array(\n            image,\n            decoded_tensors['groundtruth_boxes'].numpy(),\n            decoded_tensors['groundtruth_classes'].numpy().astype('int'),\n            scores,\n            category_index=category_index,\n            use_normalized_coordinates=use_normalized_coordinates,\n            min_score_thresh=min_score_thresh,\n            instance_masks=decoded_tensors['groundtruth_instance_masks'].numpy().astype('uint8'),\n            line_thickness=4)\n\n        plt.imshow(image)\n        plt.axis(\"off\")\n        plt.title(f\"Image-{i+1}\")\n    plt.show()\n\n\ntf_ex_decoder = TfExampleDecoder(include_mask=True)\nNUM_CLASSES = 1\ncategory_index = {1: {'id': 1, 'name': 'blood_vessel'},}\ncategory_ids = [1]\nbuffer_size = 100\nnum_of_examples = 3\n\ntrain_tfrecords = tf.io.gfile.glob(exp_config.task.train_data.input_path)\nraw_records = tf.data.TFRecordDataset(train_tfrecords).shuffle(buffer_size=buffer_size).take(num_of_examples)\nshow_batch(raw_records, num_of_examples)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with distribution_strategy.scope():\n    task = tfm.core.task_factory.get_task(exp_config.task, logging_dir=model_dir)\n    \n\nmodel, eval_logs = tfm.core.train_lib.run_experiment(\n    distribution_strategy=distribution_strategy,\n    task=task,\n    mode='train_and_eval',\n    params=exp_config,\n    model_dir=model_dir,\n    run_post_eval=True)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### export model","metadata":{}},{"cell_type":"code","source":"export_saved_model_lib.export_inference_graph(\n    input_type='image_tensor',\n    batch_size=1,\n    input_image_size=[HEIGHT, WIDTH],\n    params=exp_config,\n    checkpoint_path=tf.train.latest_checkpoint(model_dir),\n    export_dir=export_dir)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}