arnheim_3_patch_maker.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9G-3dxY01xJd"
      },
      "source": [
        "# Arnheim 3 - Segmented Patch Creator\n",
        "DeepMind, 2021\n",
        "\n",
        "## Intructions\n",
        "This Colab is to support the creation of segmented patches for collage creation using Arnheim 3.\n",
        "\n",
        "The Colab uses [PixelLib](https://github.com/ayoolaolafenwa/PixelLib) and is pretty basic but good enough to get one started creating patches from JPG images.\n",
        "\n",
        "The process is\n",
        "\n",
        "1) Provide source images\n",
        "\n",
        "* Upload images using this Colab to either Google Drive or the temp folder\n",
        "* Alternatively use a Google Drive folder that already contains images\n",
        "\n",
        "2) Create segmented patches\n",
        "* The patch file is save to Google Drive. Be sure to copy the location of the file in the Arnheim 3 Colab.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-X9f5OnKJ1-O"
      },
      "outputs": [],
      "source": [
        "#@title Installations\n",
        "!pip3 install pixellib\n",
        "!pip3 install tensorflow==2.0.1\n",
        "!pip3 install Keras==2.3.0\n",
        "!pip3 install h5py==2.10.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K41eEXVbu1k6"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "import glob\n",
        "from google.colab import drive\n",
        "from google.colab import files\n",
        "import io\n",
        "import numpy as np\n",
        "import os\n",
        "import pathlib\n",
        "import pixellib\n",
        "from pixellib.instance import instance_segmentation\n",
        "import requests\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A_k0GI-du3E1"
      },
      "outputs": [],
      "source": [
        "#@title Function definitions\n",
        "\n",
        "def mkdir(path):\n",
        "  pathlib.Path(path).mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "def upload_files(target_path):\n",
        "  \"\"\"Upload files to target directory.\"\"\"\n",
        "  mkdir(target_path)\n",
        "  uploaded = files.upload()\n",
        "  for k, v in uploaded.items():\n",
        "    open(target_path + \"/\" + k, 'wb').write(v)\n",
        "  return list(uploaded.keys())\n",
        "\n",
        "\n",
        "def download_from_url(url, force=False):\n",
        "  \"\"\"Download file from URL and cache it.\"\"\"\n",
        "\n",
        "  cache_dir = \"/content/cache\"\n",
        "  mkdir(cache_dir)\n",
        "  cache_filename = f\"{cache_dir}/{os.path.basename(url)}\"\n",
        "  cache = pathlib.Path(cache_filename)\n",
        "  if not cache.is_file() or force:\n",
        "    print(\"Downloading \" + url)\n",
        "    r = requests.get(url)\n",
        "    bytesio_object = io.BytesIO(r.content)\n",
        "    with open(cache_filename, \"wb\") as f:\n",
        "        f.write(bytesio_object.getbuffer())\n",
        "  else:\n",
        "    print(\"Using cached version of \" + url)\n",
        "  return cache "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hpsPBACa4gps"
      },
      "outputs": [],
      "source": [
        "#@title Authorise and mount Google Drive\n",
        "ROOT = \"/content\"\n",
        "MOUNT_DIR = f\"{ROOT}/drive\"\n",
        "drive.mount(MOUNT_DIR)\n",
        "# ROOT_PATH = f\"{MOUNT_DIR}/MyDrive/Arnheim3\"\n",
        "# \n",
        "# mkdir(ROOT)\n",
        "# IMAGE_PATH = f\"{ROOT}/source_images\"\n",
        "# SEGMENTED_PATH = f\"{ROOT}/segmented\"\n",
        "# \n",
        "# print(f\"\\nUsing base directory: {ROOT}\")\n",
        "# print(f\"Source images directory: {IMAGE_PATH}\")\n",
        "# print(f\"Segmented directory: {SEGMENTED_PATH}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "collapsed": true,
        "id": "f2Dhp16ptHuF"
      },
      "outputs": [],
      "source": [
        "#@title Source images and target file locations\n",
        "#@markdown Source images can be stored temporarily with the Colab, be already on Google Drive, or can be uploaded to Google Drive.\n",
        "use_google_drive_for_source_images = True   #@param {type:\"boolean\"}\n",
        "#@markdown Source images (if stored on Google Drive)\n",
        "GOOGLE_DRIVE_PATH_SOURCE_IMAGES = \"Art/Collage/Images\" #@param {type:\"string\"}\n",
        "#@markdown Target segmentation file will be saved to Google Drive for use with the Arnheim 3 Colab.\n",
        "SEGMENTED_DATA_FILENAME = \"fruit.npy\"  #@param {type: \"string\"}\n",
        "GOOGLE_DRIVE_PATH_SEGMENTED_DATA = \"Art/Collage/Patches\" #@param {type:\"string\"}\n",
        "\n",
        "data_path = MOUNT_DIR + \"/MyDrive/\" + GOOGLE_DRIVE_PATH_SEGMENTED_DATA\n",
        "data_file = data_path + \"/\" + SEGMENTED_DATA_FILENAME\n",
        "\n",
        "if use_google_drive_for_source_images:\n",
        "  IMAGE_PATH = MOUNT_DIR + \"/MyDrive/\" + GOOGLE_DRIVE_PATH_SOURCE_IMAGES\n",
        "else:\n",
        "  IMAGE_PATH = f\"{ROOT}/source_images\"\n",
        "mkdir(IMAGE_PATH)\n",
        "mkdir(data_path)\n",
        "\n",
        "print(f\"Source images directory: {IMAGE_PATH}\")\n",
        "print(f\"Segmented data will be saved to: {data_file}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Ylun-hVm5iGq"
      },
      "outputs": [],
      "source": [
        "#@title Run this cell to upload a new set of images to segment\n",
        "empty_target_dir_before_upload = False  #@param {type:\"boolean\"}\n",
        "\n",
        "if empty_target_dir_before_upload:\n",
        "  !rm {IMAGE_PATH}/*\n",
        "\n",
        "upload_files(IMAGE_PATH)\n",
        "print(f\"Images uploaded images to {IMAGE_PATH}\")\n",
        "\n",
        "!ls -l {IMAGE_PATH}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TEmqp9LOdv9P"
      },
      "outputs": [],
      "source": [
        "#@title Segment images and save patch file\n",
        "\n",
        "# https://pixellib.readthedocs.io/en/latest/Image_instance.html\n",
        "segment_image = instance_segmentation()\n",
        "segmentation_model_file = download_from_url(\n",
        "    \"https://github.com/ayoolaolafenwa/PixelLib/releases/download/1.2/mask_rcnn_coco.h5\")\n",
        "segment_image.load_model(segmentation_model_file)\n",
        "\n",
        "imagefiles = []\n",
        "for file in glob.glob(f\"{IMAGE_PATH}/*.jpg\"):\n",
        "  imagefiles.append(file)\n",
        "\n",
        "print(imagefiles)\n",
        "print(\"num images to process = \", len(imagefiles))\n",
        "\n",
        "segmented_images = []\n",
        "for imagefile in imagefiles:\n",
        "  print(imagefile)\n",
        "  try:\n",
        "    seg, _ = segment_image.segmentImage(\n",
        "        imagefile,\n",
        "        extract_segmented_objects=True,\n",
        "        save_extracted_objects =False,\n",
        "        show_bboxes=False,\n",
        "        output_image_name=str(imagefile) + \"______.tiff\")\n",
        "  except:\n",
        "    print(\"Error encounted - skipping\")\n",
        "    continue\n",
        "\n",
        "  if not len(seg[\"extracted_objects\"]):\n",
        "    print(\"Failed to segment\", imagefile)\n",
        "  else:\n",
        "    for result in seg[\"extracted_objects\"]:\n",
        "      print(result.shape)\n",
        "      segmented_image = result[..., ::-1].copy()\n",
        "      segmented_images.append(segmented_image)\n",
        "\n",
        "with open(data_file, \"wb\") as f:\n",
        "  np.save(f, segmented_images)\n",
        "print(\"Saved patch file to\", data_file)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "MakeSegmentedPatches.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}