{"id":20346653,"url":"https://github.com/arthurfdlr/pet-image-segmentation","last_synced_at":"2026-06-07T00:32:19.736Z","repository":{"id":112736952,"uuid":"362176782","full_name":"ArthurFDLR/pet-image-segmentation","owner":"ArthurFDLR","description":"🐶 Simple Jupyter Notebook to design and train your own TensorFlow model or extend an existing backbone classification model to perform pixel-wise pet image segmentation.","archived":false,"fork":false,"pushed_at":"2021-05-29T19:47:28.000Z","size":61006,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-03-04T15:48:30.862Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ArthurFDLR.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-04-27T16:18:29.000Z","updated_at":"2021-05-29T19:47:31.000Z","dependencies_parsed_at":null,"dependency_job_id":"5acc361f-ca29-44f5-8fad-1ffec2020c8f","html_url":"https://github.com/ArthurFDLR/pet-image-segmentation","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/ArthurFDLR/pet-image-segmentation","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ArthurFDLR%2Fpet-image-segmentation","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ArthurFDLR%2Fpet-image-segmentation/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ArthurFDLR%2Fpet-image-segmentation/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ArthurFDLR%2Fpet-image-segmentation/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ArthurFDLR","download_url":"https://codeload.github.com/ArthurFDLR/pet-image-segmentation/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ArthurFDLR%2Fpet-image-segmentation/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":34005030,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-06-06T02:00:07.033Z","response_time":107,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-14T22:13:44.805Z","updated_at":"2026-06-07T00:32:19.717Z","avatar_url":"https://github.com/ArthurFDLR.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003ch1 align = \"center\"\u003e Pet image segmentation \u003c/h1\u003e\n\n[![MIT License][license-shield]][license-url]\n[![LinkedIn][linkedin-shield]][linkedin-url]\n\n\u003cp align=\"center\"\u003e\n    \u003cimg src=\"./.github/banner.png\" alt=\"Generated track samples\" width=\"100%\"\u003e\n\u003c/p\u003e\n\nThe task of an image classification network is to assign a label to the input image. It is not possible to know where in the image and how big the object is. One of the solutions to tackle this task is to design a pixel-wise classification system. Each pixel of the input image is given a label. However, as always in machine learning, data is crucial. The ground truth of the dataset used to train an image segmentation neural network cannot be a simple label. It has to be a mask that matches the size of the input image. The development of such a dataset requires a lot of human labor. Fortunately, there are some open-source datasets for us to experiment with. We will use the Oxford-IIIT dataset created by Parkhi et al. It consists of images of pet and a mask where each pixel is given one of three categories; pixel belonging to the pet, pixel bordering the pet or, surrounding pixel.\n\n- [Import data](#import-data)\n- [Build models](#build-models)\n  - [Basic model](#basic-model)\n  - [Custom U-Net](#custom-u-net)\n- [Training](#training)\n- [Evaluate](#evaluate)\n\n\u003cdetails\u003e\n    \u003csummary\u003eSee package import cell\u003c/summary\u003e\n\n```python\nimport tensorflow as tf\ngpus = tf.config.experimental.list_physical_devices('GPU')\n# Limit GPU memory allocation for Conv2DTranspose execution\nif gpus:\n    try:\n    # Currently, memory growth needs to be the same across GPUs\n        for gpu in gpus:\n            tf.config.experimental.set_memory_growth(gpu, True)\n        logical_gpus = tf.config.experimental.list_logical_devices('GPU')\n        print(len(gpus), \"Physical GPUs,\", len(logical_gpus), \"Logical GPUs\")\n    except RuntimeError as e:\n    # Memory growth must be set before GPUs have been initialized\n        print(e)\n\nimport tensorflow_datasets as tfds\ntry:\n    from tensorflow_examples.models.pix2pix import pix2pix\nexcept:\n    !pip install -q git+https://github.com/tensorflow/examples.git\n    from tensorflow_examples.models.pix2pix import pix2pix\nimport numpy as np\n\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as mpatches\nfrom matplotlib.lines import Line2D\nimport matplotlib as mpl\nfrom matplotlib.ticker import MaxNLocator\nmpl.rcParams['figure.dpi'] = 300\nplt.style.use('ggplot')\n```\n\n    1 Physical GPUs, 1 Logical GPUs\n\u003c/details\u003e\n\n\u003cdetails\u003e\n    \u003csummary\u003eSee Util function definition\u003c/summary\u003e\n\n```python\ndef get_models_history(histories):\n    fig, axs = plt.subplots(1, 2, figsize=(14,5))\n    colors_graph = [\n    \"#927ced\",\n    \"#73bd4d\",\n    \"#e462c0\",\n    \"#eb5e52\"]\n    handles = []\n    for (model_name, history), color in zip(histories.items(), colors_graph):\n        axs[0].plot(range(1, len(history.history['loss'])+1), history.history['loss'], c=color, ls='-.', alpha=.7, marker='o', ms=3)\n        axs[0].plot(range(1, len(history.history['val_loss'])+1), history.history['val_loss'], label=model_name, c=color, marker='o', ms=3)\n        axs[1].plot(range(1, len(history.history['accuracy'])+1), history.history['accuracy'], c=color, ls='-.', alpha=.7, marker='o', ms=3)\n        axs[1].plot(range(1, len(history.history['val_accuracy'])+1), history.history['val_accuracy'], label=model_name, c=color, marker='o', ms=3)\n        handles.append(mpatches.Patch(color=color, label=model_name))\n    for ax in axs:\n        ax.set_xlabel('Epoch')\n        ax.xaxis.set_major_locator(MaxNLocator(integer=True))\n    axs[0].set_ylabel('loss')\n    axs[1].set_ylabel('accuracy')\n    handles.append(Line2D([0], [0], color='grey', lw=1, ls='-', label='validation'))\n    handles.append(Line2D([0], [0], color='grey', lw=1, ls='-.', label='training'))\n    fig.subplots_adjust(right=0.85)\n    fig.legend(handles=handles,\n        loc=\"center right\",\n        borderaxespad=1)\n    return fig\n\ndef display(display_list):\n    plt.figure(figsize=(15, 15))\n\n    title = ['Input Image', 'True Mask', 'UNet Prediction', 'Conv-Deconv Prediction']\n\n    for i in range(len(display_list)):\n        plt.subplot(1, len(display_list), i+1)\n        plt.title(title[i])\n        plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))\n        plt.axis('off')\n    plt.show()\n\ndef create_mask(pred_mask):\n    pred_mask = tf.argmax(pred_mask, axis=-1)\n    pred_mask = pred_mask[..., tf.newaxis]\n    return pred_mask[0]\n```\n\n\u003c/details\u003e\n\n\n## Import data\n\n\n```python\ndef normalize(input_image, input_mask):\n    input_image = tf.cast(input_image, tf.float32) / 255.0\n    input_mask -= 1\n    return input_image, input_mask\n\n@tf.function\ndef load_image_train(datapoint):\n    input_image = tf.image.resize(datapoint['image'], (128, 128))\n    input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128))\n\n    if tf.random.uniform(()) \u003e 0.5:\n        input_image = tf.image.flip_left_right(input_image)\n        input_mask = tf.image.flip_left_right(input_mask)\n\n    input_image, input_mask = normalize(input_image, input_mask)\n\n    return input_image, input_mask\n\n@tf.function\ndef load_image_test(datapoint):\n    input_image = tf.image.resize(datapoint['image'], (128, 128))\n    input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128))\n\n    input_image, input_mask = normalize(input_image, input_mask)\n\n    return input_image, input_mask\n```\n\n\n```python\ndataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)\nTRAIN_LENGTH = info.splits['train'].num_examples\nBATCH_SIZE = 32\nBUFFER_SIZE = 1000\nSTEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE\n\ntrain = dataset['train'].map(load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE)\ntest = dataset['test'].map(load_image_test)\n\ntrain_dataset = train.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()\ntrain_dataset = train_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)\ntest_dataset = test.batch(BATCH_SIZE)\n\nfor image, mask in train.take(7):\n    sample_image, sample_mask = image, mask\ndisplay([sample_image, sample_mask])\n```\n\n\n    \n![png](./.github/pet-image-segmentation_7_0.png)\n    \n\n\n## Build models\n\nWe will use an encoder-decoder convolutional architecture. The main difference with a classical CNN is the format of the representation of information. We can consider that classification CNN are simple an encoder. They extracts the key features of the images relevant to classification. Then we were able to label images given these features. Now, we want to generate a mask as output. Again, the convolutional layer generates features that contain all the relevant information to building a mask. This is the encoder. The decoder shapes the encoded information, the features, into a mask. We can use the transposed convolution (also called deconvolution) to scale up the dimensional of our data flow to match the mask size. The stride of the deconvolution defined the upscaling factor (given a dilation rate of 1). We will need a couple of deconvolution layers with a stride of 2 to compensate for the dimension reduction due to the encoder.\n\n| Layer type | Output Shape | Number of parameters |\n| -: | :-: | :-: |\n|    Conv2D | (64, 64, 16)  | 448 |\n|    Conv2D | (64, 64, 32)  | 4640 |\n|    MaxPooling2D | (32, 32, 16)  | 0 |\n|    Conv2D | (32, 32, 64)  | 18496 |\n|    Conv2DTranspose | (64,64,64) | 102464 |\n|    Conv2DTranspose | (128,128,3) | 1731 |\n\n### Basic model\n\n```python\nbasic_model = tf.keras.Sequential(\n    [\n        tf.keras.layers.Input(shape=(128,128,3)),\n        tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), strides=2, activation=\"relu\", padding=\"same\"),\n        tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), strides=1, activation=\"relu\", padding=\"same\"),\n        tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n        tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), strides=1, activation=\"relu\", padding=\"same\"),\n        tf.keras.layers.Conv2DTranspose(filters=64, kernel_size=(3, 3), strides=2, activation=\"relu\", padding=\"same\"),\n        tf.keras.layers.Conv2DTranspose(filters=3, kernel_size=(3, 3), strides=2, padding=\"same\")\n    ]\n)\nbasic_model.summary()\n```\n\u003cdetails\u003e\n    \u003csummary\u003eSee model resume\u003c/summary\u003e\n\n    Model: \"sequential\"\n    _________________________________________________________________\n    Layer (type)                 Output Shape              Param #   \n    =================================================================\n    conv2d (Conv2D)              (None, 64, 64, 16)        448       \n    _________________________________________________________________\n    conv2d_1 (Conv2D)            (None, 64, 64, 32)        4640      \n    _________________________________________________________________\n    max_pooling2d (MaxPooling2D) (None, 32, 32, 32)        0         \n    _________________________________________________________________\n    conv2d_2 (Conv2D)            (None, 32, 32, 64)        18496     \n    _________________________________________________________________\n    conv2d_transpose (Conv2DTran (None, 64, 64, 64)        36928     \n    _________________________________________________________________\n    conv2d_transpose_1 (Conv2DTr (None, 128, 128, 3)       1731      \n    =================================================================\n    Total params: 62,243\n    Trainable params: 62,243\n    Non-trainable params: 0\n    _________________________________________________________________\n\n\u003c/details\u003e\n    \n### Custom U-Net\n\n\n```python\nencoder_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False)\n\n# Use the activations of these layers\nlayer_names = [\n    'block_1_expand_relu',   # 64x64\n    'block_3_expand_relu',   # 32x32\n    'block_6_expand_relu',   # 16x16\n    'block_13_expand_relu',  # 8x8\n    'block_16_project',      # 4x4\n]\nencoder_model_outputs = [encoder_model.get_layer(name).output for name in layer_names]\n\n# Create the feature extraction model\nencoder_stack = tf.keras.Model(inputs=encoder_model.input, outputs=encoder_model_outputs)\n\nencoder_stack.trainable = False\n\ndecoder_stack = [\n    pix2pix.upsample(512, 3),  # 4x4 -\u003e 8x8\n    pix2pix.upsample(256, 3),  # 8x8 -\u003e 16x16\n    pix2pix.upsample(128, 3),  # 16x16 -\u003e 32x32\n    pix2pix.upsample(64, 3),   # 32x32 -\u003e 64x64\n]\n\n\ninputs = tf.keras.layers.Input(shape=[128, 128, 3])\n# Downsampling through the model\nskips = encoder_stack(inputs)\nx = skips[-1]\nskips = reversed(skips[:-1])\n\n# Upsampling and establishing the skip connections\nfor up, skip in zip(decoder_stack, skips):\n    x = up(x)\n    concat = tf.keras.layers.Concatenate()\n    x = concat([x, skip])\n\n# This is the last layer of the model\nlast = tf.keras.layers.Conv2DTranspose(\n  3, 3, strides=2,\n  padding='same')  #64x64 -\u003e 128x128\n\nx = last(x)\n\nUNet_model = tf.keras.Model(inputs=inputs, outputs=x)\nUNet_model.summary()\n```\n\n\u003cdetails\u003e\n    \u003csummary\u003eSee model resume\u003c/summary\u003e\n\n    Model: \"functional_3\"\n    __________________________________________________________________________________________________\n    Layer (type)                    Output Shape         Param #     Connected to                     \n    ==================================================================================================\n    input_3 (InputLayer)            [(None, 128, 128, 3) 0                                            \n    __________________________________________________________________________________________________\n    functional_1 (Functional)       [(None, 64, 64, 96), 1841984     input_3[0][0]                    \n    __________________________________________________________________________________________________\n    sequential_1 (Sequential)       (None, 8, 8, 512)    1476608     functional_1[0][4]               \n    __________________________________________________________________________________________________\n    concatenate (Concatenate)       (None, 8, 8, 1088)   0           sequential_1[0][0]               \n                                                                     functional_1[0][3]               \n    __________________________________________________________________________________________________\n    sequential_2 (Sequential)       (None, 16, 16, 256)  2507776     concatenate[0][0]                \n    __________________________________________________________________________________________________\n    concatenate_1 (Concatenate)     (None, 16, 16, 448)  0           sequential_2[0][0]               \n                                                                     functional_1[0][2]               \n    __________________________________________________________________________________________________\n    sequential_3 (Sequential)       (None, 32, 32, 128)  516608      concatenate_1[0][0]              \n    __________________________________________________________________________________________________\n    concatenate_2 (Concatenate)     (None, 32, 32, 272)  0           sequential_3[0][0]               \n                                                                     functional_1[0][1]               \n    __________________________________________________________________________________________________\n    sequential_4 (Sequential)       (None, 64, 64, 64)   156928      concatenate_2[0][0]              \n    __________________________________________________________________________________________________\n    concatenate_3 (Concatenate)     (None, 64, 64, 160)  0           sequential_4[0][0]               \n                                                                     functional_1[0][0]               \n    __________________________________________________________________________________________________\n    conv2d_transpose_6 (Conv2DTrans (None, 128, 128, 3)  4323        concatenate_3[0][0]              \n    ==================================================================================================\n    Total params: 6,504,227\n    Trainable params: 4,660,323\n    Non-trainable params: 1,843,904\n    __________________________________________________________________________________________________\n\n\u003c/details\u003e\n\n## Training\n\nNote that the output has three channels, whereas provided mask ground truth only has one. Indeed, we use one-hot encoding for the output of the model to ease training and computation. The difference of dimension is taken into account into the calculation of the loss: _SparseCategoricalCrossentropy_ from logits (and not a probability distribution). Also, the same training methods can be implemented through the one-hot encoding of the ground-truth and _CategoricalCrossentropy_ loss function.\n\nWe will also train a modified U-Net model in addition to the low-complexity architecture proposed earlier. The backbone model is a pre-trained MobilNetV2, and the pix2pix blocks are used as the decoder. This encoder-decoder model reaches a complexity of 6 504 227 parameters, which is 50 times more than the convolution-deconvolution network. Such a model requires a longer training process but yields far better results.\n\n```python\nEPOCHS = 25\nVAL_SUBSPLITS = 5\nVALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS\n\ntraining_histories = {}\n```\n\n\n```python\nname = 'Conv-Deconv'\nweights_path = \".\\weights\\{}.h5\".format(name)\ncheckpointer = tf.keras.callbacks.ModelCheckpoint(filepath=weights_path, verbose=0, save_best_only=True)\n\nbasic_model.compile(optimizer='adam',\n              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n              metrics=['accuracy'])\ntraining_histories[name] = basic_model.fit(train_dataset, epochs=EPOCHS,\n                          steps_per_epoch=STEPS_PER_EPOCH,\n                          validation_steps=VALIDATION_STEPS,\n                          validation_data=test_dataset,\n                          shuffle=True,\n                          callbacks=[checkpointer],\n                          verbose=0)\n```\n\n\n```python\nname = 'U-Net'\nweights_path = \".\\weights\\{}.h5\".format(name)\ncheckpointer = tf.keras.callbacks.ModelCheckpoint(filepath=weights_path, verbose=0, save_best_only=True)\n\nUNet_model.compile(optimizer='adam',\n              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n              metrics=['accuracy'])\ntraining_histories[name] = UNet_model.fit(train_dataset, epochs=EPOCHS,\n                          steps_per_epoch=STEPS_PER_EPOCH,\n                          validation_steps=VALIDATION_STEPS,\n                          validation_data=test_dataset,\n                          shuffle=True,\n                          callbacks=[checkpointer],\n                          verbose=0)\n```\n\n\n```python\nget_models_history(training_histories).show()\n```\n    \n![png](./.github/pet-image-segmentation_17_1.png)\n\nOnly the best-performing set of parameters is saved as the final model. The training of the custom U-Net turns out to require way fewer epochs than the Conv-Deconv model. The final model implements the set of parameters produced at the end of the sixth epochs. We can see significant overfitting in the validation loss curve afterward. Besides, some output images with semantic segmentation labels of both models are presented in the next section. As expected, the more complex model yields better results.\n\n## Evaluate\n\n\n```python\nconv_deconv_model = tf.keras.models.load_model(\"./weights/Conv-Deconv.h5\")\nUNet_model = tf.keras.models.load_model(\"./weights/U-Net.h5\")\n\nfor image, mask in test_dataset.take(5):\n    UNet_mask = UNet_model.predict(image)\n    conv_deconv_mask = conv_deconv_model.predict(image)\n    display([image[0], mask[0], create_mask(UNet_mask), create_mask(conv_deconv_mask)])\n```\n\n![png](./.github/PetIIIT_Examples.png)\n\n\u003c!-- MARKDOWN LINKS \u0026 IMAGES --\u003e\n[license-shield]: https://img.shields.io/github/license/ArthurFDLR/pet-image-segmentation?style=for-the-badge\n[license-url]: https://github.com/ArthurFDLR/pet-image-segmentation/blob/master/LICENSE\n[linkedin-shield]: https://img.shields.io/badge/-LinkedIn-black.svg?style=for-the-badge\u0026logo=linkedin\u0026colorB=555\n[linkedin-url]: https://linkedin.com/in/arthurfdlr/","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Farthurfdlr%2Fpet-image-segmentation","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Farthurfdlr%2Fpet-image-segmentation","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Farthurfdlr%2Fpet-image-segmentation/lists"}