01_standalone_examples/image-segmentation/ImageSegmentationSetup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "43e3aa9c-9a20-4468-9846-4c6f5902b7c4",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "<img src=\"https://lakefs.io/wp-content/uploads/2022/09/lakeFS-Logo.svg\" alt=\"lakeFS logo\" width=200/>\n",
    "\n",
    "# ML Data Version Control and Reproducibility at Scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "d9b7d4ff-2a76-40a4-8629-8818f667bfb8",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "source": [
    "## Setup\n",
    "\n",
    "**(you shouldn't need to change anything in this section, just run it)**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "c785bf65-a97a-4c81-8120-555028db4a0c",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "77b90199-50aa-4bc9-8183-6cc6759f8cc4",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%xmode Minimal\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import logging\n",
    "logging.basicConfig()\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "\n",
    "import lakefs\n",
    "from lakefs.exceptions import NotFoundException\n",
    "import os\n",
    "\n",
    "import boto3\n",
    "import random\n",
    "import time\n",
    "import datetime\n",
    "from pyspark.sql.functions import substring_index, col, pandas_udf, collect_list, size, desc, base64\n",
    "from PIL import Image\n",
    "import io\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from petastorm.spark import SparkDatasetConverter, make_spark_converter\n",
    "from petastorm import TransformSpec\n",
    "from functools import partial\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "import pytorch_lightning as pl\n",
    "import segmentation_models_pytorch as smp\n",
    "import base64\n",
    "import mlflow\n",
    "from hyperopt import fmin, rand, hp, SparkTrials, STATUS_OK, space_eval, tpe\n",
    "import gc\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import Markdown as md"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "412feb1b-6aa5-42de-a885-04da44ed9fa5",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Versioning Information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6d85e5ed-a8e3-45af-828e-88bc99cb3d5c",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "mainBranch = \"main\"\n",
    "emptyBranch = \"empty\"\n",
    "experimentBranch = localOrDistributedComputing + \"-experiment\"\n",
    "commitMetadata=\"\"\n",
    "tagPrefix = datetime.datetime.now().strftime(\"%Y_%m_%d_%H_%M_%S\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"LAKECTL_SERVER_ENDPOINT_URL\"] = lakefsEndPoint\n",
    "os.environ[\"LAKECTL_CREDENTIALS_ACCESS_KEY_ID\"] = lakefsAccessKey\n",
    "os.environ[\"LAKECTL_CREDENTIALS_SECRET_ACCESS_KEY\"] = lakefsSecretKey"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define lakeFS UI Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if lakefsEndPoint.startswith('http://host.docker.internal'):\n",
    "    lakefsUIEndPoint = 'http://127.0.0.1:8000'\n",
    "else:\n",
    "    lakefsUIEndPoint = lakefsEndPoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Verify lakeFS credentials by getting lakeFS version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "172f81be-53c1-4c42-bbfc-ab3db4a39b2a",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Verifying lakeFS credentials…\")\n",
    "try:\n",
    "    v=lakefs.client.Client().version\n",
    "except:\n",
    "    print(\"🛑 failed to get lakeFS version\")\n",
    "else:\n",
    "    print(f\"…✅lakeFS credentials verified\\n\\nℹ️lakeFS version {v}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "ecc978df-d338-409a-8509-8c2142ebd298",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Define lakeFS Repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "3dc99a61-962d-45e9-8c0c-a6ab96ba4745",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "repo = lakefs.Repository(repo_name).create(storage_namespace=f\"{storageNamespace}/{repo_name}\", default_branch=mainBranch, exist_ok=True)\n",
    "branchMain = repo.branch(mainBranch)\n",
    "print(repo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "3a9f2712-23f7-4ba5-b47b-6f7ae4202da5",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Configure lakectl for LOCAL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "75c61607-f816-4aa3-9ec1-8107104f8d54",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    f = open(\".lakectl.yaml\", \"w\")\n",
    "    f.write(f\"credentials: \\n\\\n",
    "    access_key_id: {lakefsAccessKey}\\n\\\n",
    "    secret_access_key: {lakefsSecretKey}\\n\\\n",
    "server: \\n\\\n",
    "    endpoint_url: {lakefsEndPoint}\")\n",
    "    f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "01aca111-3337-4545-a209-29d6a0a13593",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Set up Spark for LOCAL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "01cc56b3-f532-4c5e-9a0c-6e10f7137d72",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if localOrDistributedComputing == \"LOCAL\":\n",
    "    from pyspark.sql import SparkSession\n",
    "    spark = SparkSession.builder.appName(\"lakeFS / Jupyter\") \\\n",
    "        .config(\"spark.hadoop.fs.s3.impl\", \"org.apache.hadoop.fs.s3a.S3AFileSystem\") \\\n",
    "        .config(\"spark.hadoop.fs.s3a.endpoint\", lakefsEndPoint) \\\n",
    "        .config(\"spark.hadoop.fs.s3a.path.style.access\", \"true\") \\\n",
    "        .config(\"spark.hadoop.fs.s3a.access.key\", lakefsAccessKey) \\\n",
    "        .config(\"spark.hadoop.fs.s3a.secret.key\", lakefsSecretKey) \\\n",
    "        .config(\"spark.jars.packages\", \"io.delta:delta-core_2.12:2.3.0\") \\\n",
    "        .config(\"spark.sql.extensions\", \"io.delta.sql.DeltaSparkSessionExtension\") \\\n",
    "        .config(\"spark.sql.catalog.spark_catalog\", \"org.apache.spark.sql.delta.catalog.DeltaCatalog\") \\\n",
    "        .getOrCreate()\n",
    "    \n",
    "#        .config(\"spark.jars.packages\", \"io.delta:delta-core_2.12:2.3.0,org.mlflow.mlflow-spark\") \\\n",
    "    spark.sparkContext.setLogLevel(\"INFO\")\n",
    "\n",
    "    spark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fcef1ef4-d2e5-4f87-b4ec-370906cc895e",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Create empty branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "dc14b3a6-2b3b-4701-a466-cf0f476e4446",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "branchEmpty = repo.branch(emptyBranch).create(source_reference=mainBranch, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "cf987466-9413-453d-9675-0e55ab633948",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Folder structure to implement Medallion Architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fbda8000-e809-4e2d-9fa1-c58710cedab4",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "raw_data_folder = \"raw/\"\n",
    "bronze_data_folder = \"bronze/\"\n",
    "silver_data_folder = \"silver/\"\n",
    "gold_data_folder = \"gold/\"\n",
    "training_data_folder = \"train_v2/\"\n",
    "mask_data_folder = \"mask/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "85126538-183e-4918-b647-07211cf6a27b",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Create S3 client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "5b6649ff-7742-4f4a-903d-efcc0628c51d",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "s3 = boto3.client('s3',\n",
    "    endpoint_url='https://s3.' + awsRegion + '.amazonaws.com',\n",
    "    aws_access_key_id=aws_access_key_id,\n",
    "    aws_secret_access_key=aws_secret_access_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e5e8bd4b-5160-44e3-a4d9-aa332cef1edb",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Get the list of images in training dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1a646985-b578-4a6c-ad47-faf4034842af",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def list_images():\n",
    "    s3_result =  s3.list_objects_v2(Bucket=bucketName, Prefix=prefix+training_data_folder)\n",
    "\n",
    "    file_list = []\n",
    "    for key in s3_result['Contents']:\n",
    "        file_list.append(key['Key'])\n",
    "\n",
    "    while s3_result['IsTruncated']:\n",
    "        continuation_key = s3_result['NextContinuationToken']\n",
    "        s3_result = s3.list_objects_v2(Bucket=bucketName, Prefix=prefix, ContinuationToken=continuation_key)\n",
    "        for key in s3_result['Contents']:\n",
    "            file_list.append(key['Key'])\n",
    "    print(f\"Image count = {len(file_list)}\")\n",
    "    return file_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "52b5fcc4-0aa1-48d2-a80b-821df7eb8fcd",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Import subset of training data to lakeFS repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "a477094e-45ba-476f-926e-86fb79e65349",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def import_images(file_list_random):\n",
    "    importer = branchExperimentBranchN.import_data(commit_message=\"import images\")\n",
    "    for file in file_list_random:\n",
    "        importer.prefix(\"s3://\"+bucketName+'/'+file, destination=raw_data_folder+training_data_folder)\n",
    "\n",
    "    importer.prefix(\"s3://\"+bucketName+'/'+prefix+'train_ship_segmentations_v2.csv', destination=raw_data_folder)\n",
    "    \n",
    "    importer.start()\n",
    "    time.sleep(2)\n",
    "    status = importer.status()\n",
    "    print(status)\n",
    "\n",
    "    while not status.completed and status.error is None:\n",
    "        time.sleep(2)\n",
    "        status = importer.status()\n",
    "        print(status)\n",
    "\n",
    "    if status.error:\n",
    "        raise Exception(status.error)\n",
    "    \n",
    "    print(f\"\\nImported a total of {status.ingested_objects} objects into branch {experimentBranchN}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "62d1893d-9d3a-4901-859d-137f55465ffa",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "## Build the data pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "95dea8a4-ea85-4e56-9e80-c943822b8f8d",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Ingest raw images as bronze data set and save as Delta table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e5867767-45b9-4454-ac76-aa14fc388f26",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def bronze_images():\n",
    "    return spark.read.format(\"binaryfile\").option(\"pathGlobFilter\", \"*.jpg\").load(training_data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "219f4da1-6b3d-4a39-a478-949ed904857a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Print Diff results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "9fa59cb8-9ef1-4573-91c6-bea3fc3d2cb2",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def print_diff_refs(diff_refs):\n",
    "    results = map(\n",
    "        lambda n:[n.path,n.path_type,n.size_bytes,n.type],\n",
    "        diff_refs)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0f62b5f2-0323-478f-8346-8a69a8be878e",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Diff the branch to find uncommitted changes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "c75ad558-3866-4b86-bb9a-a15a4fba5b53",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def diff_branch(repo, repo_path, branch):\n",
    "    if localOrDistributedComputing == \"LOCAL\":\n",
    "        lakeFSLocalCommand = f\"lakectl local status {repo_path}\"\n",
    "        ! $lakeFSLocalCommand\n",
    "    elif localOrDistributedComputing == \"DISTRIBUTED\":\n",
    "        for a in print_diff_refs(\n",
    "            branchMain.diff(other_ref=branch)):\n",
    "            print(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "96dd6218-ea87-49cf-a2ed-492b648fd79a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Commit changes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "b112b57c-f56f-4278-a065-226479355595",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def print_lakectl_response(response, lines):\n",
    "    if lines > len(response):\n",
    "        lines = len(response)\n",
    "    for x in range(lines):\n",
    "        print(response[-(lines-x)])\n",
    "\n",
    "def get_commit_id(commit_response):\n",
    "    responseLastNLinesSplit = [i.split(':') for i in commit_response]\n",
    "    for key_value in responseLastNLinesSplit:\n",
    "        for key in key_value:\n",
    "            if key == 'ID':\n",
    "                return key_value[1].strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "7f784f4a-3ed5-49de-98da-2afd4001e34c",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def commit(repo, repo_path, branch, commitMessage, metadata=\"\"):\n",
    "    if localOrDistributedComputing == \"LOCAL\":\n",
    "        lakeFSLocalCommand = f\"lakectl local commit -m '{commitMessage}' --meta '{metadata}' {repo_path}\"\n",
    "        commit_response = ! $lakeFSLocalCommand\n",
    "        # Print last 15 lines\n",
    "        print_lakectl_response(commit_response, 15)\n",
    "        return get_commit_id(commit_response)\n",
    "    elif localOrDistributedComputing == \"DISTRIBUTED\":\n",
    "        if metadata == \"\":\n",
    "            metadata = {}\n",
    "        ref = lakefs.Repository(repo).branch(branch).commit(message=commitMessage, metadata=metadata)\n",
    "        return ref.get_commit().id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "5991ef81-03af-4e4a-b117-4d0f6a7c70d9",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Create lakeFS tag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "8e47ac94-9317-42d5-8338-2c947cb21d89",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def lakefs_set_tag(repo, tagID, branch):\n",
    "    print(lakefs.Tag(repo, tagID).create(branch, exist_ok=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "9a7d4dd7-5581-4a12-a14e-6a959c2910cc",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Enrich dataset and save as silver dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b0aded32-476b-41a0-abba-fe27c695402f",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "IMAGE_RESIZE = 320 # divisible by 32\n",
    "\n",
    "# Resize UDF function\n",
    "@pandas_udf(\"binary\")\n",
    "def resize_image_udf(content_series):\n",
    "    def resize_image(content):\n",
    "        \"\"\"resize image and serialize as jpeg\"\"\"\n",
    "        try:\n",
    "            image = Image.open(io.BytesIO(content)).resize((IMAGE_RESIZE, IMAGE_RESIZE), Image.NEAREST)\n",
    "            output = io.BytesIO()\n",
    "            image.save(output, format='JPEG')\n",
    "            return output.getvalue()\n",
    "        except Exception:\n",
    "            # some images are invalid\n",
    "            return None\n",
    "\n",
    "    return content_series.apply(resize_image)\n",
    "\n",
    "# add the metadata to enable the image preview\n",
    "image_meta = {\"spark.contentAnnotation\": '{\"mimeType\": \"image/jpeg\"}'}\n",
    "\n",
    "def silver_images(df_bronze_images):\n",
    "    return df_bronze_images.withColumn(\"image_id\", substring_index(col('path'), '/', -1)) \\\n",
    "        .withColumn(\"content\", resize_image_udf(col(\"content\")).alias(\"content\", metadata=image_meta)) \\\n",
    "        .filter(\"content is not null\") \\\n",
    "        .select(\"image_id\", \"content\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "851ccf08-2446-44dd-b448-bd7a992d5982",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Load the raw image mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b4109119-51be-41c2-9b4b-9ca1b6869f70",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def bronze_mask():\n",
    "    annotationsDF = (spark.read.option(\"header\", \"true\")\n",
    "                     .option(\"inferSchema\", \"true\")\n",
    "                     .csv(f\"{raw_data_path}/train_ship_segmentations_v2.csv\"))\n",
    "    return (annotationsDF.withColumnRenamed(\"ImageId\", \"image_id\")\n",
    "            .withColumnRenamed(\"EncodedPixels\", \"encoded_pixels\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1e6c26b1-8895-4782-81f1-5864c5098d9a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Transforming masks into images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "5b5894b8-1a57-414b-a5ad-bcb13ae116dc",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "IMAGE_SIZE = 768\n",
    "\n",
    "def rle_decode(mask_rle, shape=(IMAGE_SIZE, IMAGE_SIZE)):\n",
    "    s = mask_rle.split()\n",
    "    starts = np.asarray(s[0::2], dtype=int)\n",
    "    lengths = np.asarray(s[1::2], dtype=int)\n",
    "\n",
    "    starts -= 1\n",
    "    ends = starts + lengths\n",
    "    img = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n",
    "    for lo, hi in zip(starts, ends):\n",
    "        img[lo:hi] = 255\n",
    "    return img.reshape(shape).T\n",
    "\n",
    "\n",
    "def mask(rle_masks):\n",
    "    if isinstance(rle_masks, np.ndarray):\n",
    "        all_masks = np.zeros((IMAGE_SIZE, IMAGE_SIZE), dtype=np.int8)\n",
    "        for mask in rle_masks.tolist():\n",
    "            all_masks += rle_decode(mask)\n",
    "        image = Image.fromarray(all_masks, mode=\"L\").resize((IMAGE_RESIZE, IMAGE_RESIZE), Image.NEAREST)\n",
    "        output = io.BytesIO()\n",
    "        image.save(output, format='JPEG')\n",
    "        return output.getvalue()\n",
    "    raise Exception(type(rle_masks))\n",
    "\n",
    "\n",
    "@pandas_udf(\"binary\")\n",
    "def computeMaskUDF(s: pd.Series) -> pd.Series:\n",
    "    return s.apply(mask)\n",
    "\n",
    "\n",
    "def silver_mask(df_bronze_mask):\n",
    "    return (df_bronze_mask\n",
    "            .filter(\"encoded_pixels is not null\")\n",
    "            .groupBy(\"image_id\").agg(collect_list('encoded_pixels').alias('encoded_pixels'))\n",
    "            .withColumn(\"boat_number\", size(col(\"encoded_pixels\")))\n",
    "            .withColumn(\"mask\", computeMaskUDF(col(\"encoded_pixels\")).alias(\"mask\", metadata=image_meta)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6b86d14d-bf0d-44e5-9daa-9408cbac586c",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Joining image and mask both as the gold layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ba4f442a-bace-4722-a537-897e82cfaba8",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def gold_images(df_silver_images, df_silver_mask):\n",
    "    imagesWithMask = df_silver_mask.join(df_silver_images, \"image_id\")\n",
    "    return imagesWithMask.select(\"image_id\", \"boat_number\", \"mask\", \"content\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "5ff7e332-a275-47ea-a123-e3483f329087",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Decode the raw image bytes and apply standard ImageNet transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "dc7b3f6d-7615-440a-b285-6e6cc75295a4",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def transform_row(is_train, pd_batch):\n",
    "  \"\"\"\n",
    "  The input and output of this function must be pandas dataframes.\n",
    "  Do data augmentation for the training dataset only.\n",
    "  \"\"\"\n",
    "  transformers = [transforms.Lambda(lambda x: Image.open(io.BytesIO(x)))]\n",
    "  if is_train:\n",
    "    transformers.extend([\n",
    "      transforms.RandomResizedCrop(224),\n",
    "      transforms.RandomHorizontalFlip(),\n",
    "    ])\n",
    "  else:\n",
    "    transformers.extend([\n",
    "      transforms.Resize(256),\n",
    "      transforms.CenterCrop(224),\n",
    "    ])\n",
    "  transformers.extend([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
    "  ])\n",
    "  \n",
    "  transformersMask = [transforms.Lambda(lambda x: Image.open(io.BytesIO(x)))]\n",
    "  if is_train:\n",
    "    transformersMask.extend([\n",
    "      #transforms.RandomResizedCrop(224),\n",
    "      #transforms.RandomHorizontalFlip(),\n",
    "      transforms.Resize(256),\n",
    "      transforms.CenterCrop(224),\n",
    "    ])\n",
    "  else:\n",
    "    transformersMask.extend([\n",
    "      transforms.Resize(256),\n",
    "      transforms.CenterCrop(224),\n",
    "    ])\n",
    "  transformersMask.extend([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5]),\n",
    "  ])\n",
    "  \n",
    "  trans = transforms.Compose(transformers)\n",
    "  transMask = transforms.Compose(transformersMask)\n",
    "\n",
    "  pd_batch['content'] = pd_batch['content'].map(lambda x: trans(x).numpy())\n",
    "  pd_batch['mask'] = pd_batch['mask'].map(lambda x: transMask(x).numpy())\n",
    "  #pd_batch = pd_batch.drop(labels=['content'], axis=1)\n",
    "  pd_batch = pd_batch.drop(labels=['image_id'], axis=1)\n",
    "  pd_batch = pd_batch.drop(labels=['boat_number'], axis=1)\n",
    "  #pd_batch = pd_batch.drop(labels=['mask'], axis=1)\n",
    "  return pd_batch\n",
    "\n",
    "def get_transform_spec(is_train=True):\n",
    "  # Note that the output shape of the `TransformSpec` is not automatically known by petastorm, \n",
    "  # so we need to specify the shape for new columns in `edit_fields` and specify the order of \n",
    "  # the output columns in `selected_fields`.\n",
    "  return TransformSpec(partial(transform_row, is_train), \n",
    "                       edit_fields=[('content', np.float32, (3, 224, 224), False), ('mask', np.float32, (1, 224, 224), False)], # num_channels is 1 for greyscale images\n",
    "                       #edit_fields=[('content', torch.bfloat16, (3, 224, 224), False), ('mask', torch.bfloat16, (1, 224, 224), False)], \n",
    "                       selected_fields=['content', 'mask'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "61ccfd3a-743a-40cd-825e-d45a333d2760",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7e7ddf75-0583-4718-86e0-aa20db4086d9",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class BoatModel(pl.LightningModule):\n",
    "\n",
    "    def __init__(self, arch, encoder_name, in_channels, out_classes, lr, **kwargs):\n",
    "        super().__init__()\n",
    "        self.model = smp.create_model(\n",
    "            arch, encoder_name=encoder_name, in_channels=in_channels, classes=out_classes, **kwargs\n",
    "        )\n",
    "\n",
    "        # preprocessing parameteres for image\n",
    "        params = smp.encoders.get_preprocessing_params(encoder_name)\n",
    "        self.register_buffer(\"std\", torch.tensor(params[\"std\"]).view(1, 3, 1, 1))\n",
    "        self.register_buffer(\"mean\", torch.tensor(params[\"mean\"]).view(1, 3, 1, 1))\n",
    "\n",
    "        # for image segmentation dice loss could be the best first choice\n",
    "        self.loss_fn = smp.losses.DiceLoss(smp.losses.BINARY_MODE, from_logits=True)\n",
    "\n",
    "    def forward(self, image):\n",
    "        # normalize image here\n",
    "        image = (image - self.mean) / self.std\n",
    "        mask = self.model(image)\n",
    "        return mask\n",
    "\n",
    "    def shared_step(self, batch, stage):\n",
    "        \n",
    "        image = batch[\"content\"]\n",
    "\n",
    "        # Shape of the image should be (batch_size, num_channels, height, width)\n",
    "        # if you work with grayscale images, expand channels dim to have [batch_size, 1, height, width]\n",
    "        assert image.ndim == 4\n",
    "\n",
    "        # Check that image dimensions are divisible by 32, \n",
    "        # encoder and decoder connected by `skip connections` and usually encoder have 5 stages of \n",
    "        # downsampling by factor 2 (2 ^ 5 = 32); e.g. if we have image with shape 65x65 we will have \n",
    "        # following shapes of features in encoder and decoder: 84, 42, 21, 10, 5 -> 5, 10, 20, 40, 80\n",
    "        # and we will get an error trying to concat these features\n",
    "        h, w = image.shape[2:]\n",
    "        assert h % 32 == 0 and w % 32 == 0\n",
    "\n",
    "        mask = batch[\"mask\"]\n",
    "\n",
    "        # Shape of the mask should be [batch_size, num_classes, height, width]\n",
    "        # for binary segmentation num_classes = 1\n",
    "        assert mask.ndim == 4\n",
    "\n",
    "        # Check that mask values in between 0 and 1, NOT 0 and 255 for binary segmentation\n",
    "        #assert mask.max() <= 1.0 and mask.min() >= 0 # Commented to avoid error\n",
    "\n",
    "        logits_mask = self.forward(image)\n",
    "        \n",
    "        # Predicted mask contains logits, and loss_fn param `from_logits` is set to True\n",
    "        loss = self.loss_fn(logits_mask, mask)\n",
    "\n",
    "        # Lets compute metrics for some threshold\n",
    "        # first convert mask values to probabilities, then \n",
    "        # apply thresholding\n",
    "        prob_mask = logits_mask.sigmoid()\n",
    "        pred_mask = (prob_mask > 0.5).float()\n",
    "\n",
    "        # We will compute IoU metric by two ways\n",
    "        #   1. dataset-wise\n",
    "        #   2. image-wise\n",
    "        # but for now we just compute true positive, false positive, false negative and\n",
    "        # true negative 'pixels' for each image and class\n",
    "        # these values will be aggregated in the end of an epoch\n",
    "        tp, fp, fn, tn = smp.metrics.get_stats(pred_mask.long(), mask.long(), mode=\"binary\")\n",
    "\n",
    "        return {\n",
    "            \"loss\": loss,\n",
    "            \"tp\": tp,\n",
    "            \"fp\": fp,\n",
    "            \"fn\": fn,\n",
    "            \"tn\": tn,\n",
    "        }\n",
    "\n",
    "    def shared_epoch_end(self, outputs, stage):\n",
    "        # aggregate step metics\n",
    "        tp = torch.cat([x[\"tp\"] for x in outputs])\n",
    "        fp = torch.cat([x[\"fp\"] for x in outputs])\n",
    "        fn = torch.cat([x[\"fn\"] for x in outputs])\n",
    "        tn = torch.cat([x[\"tn\"] for x in outputs])\n",
    "        \n",
    "        # per image IoU means that we first calculate IoU score for each image \n",
    "        # and then compute mean over these scores\n",
    "        per_image_iou = smp.metrics.iou_score(tp, fp, fn, tn, reduction=\"micro-imagewise\")\n",
    "        \n",
    "        # dataset IoU means that we aggregate intersection and union over whole dataset\n",
    "        # and then compute IoU score. The difference between dataset_iou and per_image_iou scores\n",
    "        # in this particular case will not be much, however for dataset \n",
    "        # with \"empty\" images (images without target class) a large gap could be observed. \n",
    "        # Empty images influence a lot on per_image_iou and much less on dataset_iou.\n",
    "        dataset_iou = smp.metrics.iou_score(tp, fp, fn, tn, reduction=\"micro\")\n",
    "\n",
    "        metrics = {\n",
    "            f\"{stage}_per_image_iou\": per_image_iou,\n",
    "            f\"{stage}_dataset_iou\": dataset_iou,\n",
    "        }\n",
    "        \n",
    "        self.log_dict(metrics, prog_bar=True)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        return self.shared_step(batch, \"train\")            \n",
    "\n",
    "    def training_epoch_end(self, outputs):\n",
    "        return self.shared_epoch_end(outputs, \"train\")\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        return self.shared_step(batch, \"valid\")\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        return self.shared_epoch_end(outputs, \"valid\")\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        return self.shared_step(batch, \"test\")  \n",
    "\n",
    "    def test_epoch_end(self, outputs):\n",
    "        return self.shared_epoch_end(outputs, \"test\")\n",
    "\n",
    "    def predict_step(self, batch, batch_idx):\n",
    "        return self.shared_step(batch, \"test\")  \n",
    "\n",
    "    def predict_epoch_end(self, outputs):\n",
    "        return self.shared_epoch_end(outputs, \"test\")\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam(self.parameters(), lr=0.0001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "9883528c-0446-4be7-8e1a-94bfe35c1c06",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "###  Model wrapper which tackles all the data transformations originally captured in transform spec."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b7781d73-34fb-4772-b26c-ad4231659325",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class CVModelWrapper(mlflow.pyfunc.PythonModel):\n",
    "  \n",
    "  def __init__(self, model):    \n",
    "    # instantiate model in evaluation mode\n",
    "    self.model = model.eval()\n",
    "    \n",
    "     # define transformation pipeline\n",
    "    trans = torchvision.transforms.Compose([\n",
    "              torchvision.transforms.Lambda(lambda x: Image.open(io.BytesIO(x))),\n",
    "              torchvision.transforms.Resize(256),\n",
    "              torchvision.transforms.ToTensor(),\n",
    "              torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])\n",
    "              ])\n",
    "    self.transform = trans\n",
    "    \n",
    "  def predict(self, context, model_input):\n",
    "    for i in torch.utils.data.DataLoader(model_input):\n",
    "      output = model(i)\n",
    "      #probs = torch.nn.functional.softmax(output, dim=1)[:,1]\n",
    "      outputs += [output]\n",
    "    \n",
    "    ser = pd.Series(outputs)\n",
    "    return ser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "37ca43f7-4343-4096-ac15-68923835fe73",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Train the base Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2fdd782a-5370-4e76-9a8a-e5cdc0557fb9",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16\n",
    "\n",
    "#mlflow.autolog(log_models=False)\n",
    "\n",
    "def train_model(arch, encoder_name, lr, nested=False):\n",
    "    with mlflow.start_run(nested=nested) as run:\n",
    "        model = BoatModel(arch=arch, encoder_name=encoder_name, in_channels=3, out_classes=1, lr=lr)\n",
    "        trainer = pl.Trainer(\n",
    "            benchmark=True,\n",
    "            num_sanity_val_steps=0,\n",
    "            precision=32, # Changed precision from 16 to 32 to run model on CPUs\n",
    "            #accelerator=\"gpu\",\n",
    "            #gpus=\"1\",\n",
    "            log_every_n_steps=100,\n",
    "            default_root_dir=\"/tmp\",\n",
    "            max_epochs=1)\n",
    "\n",
    "        mlflow.log_param(\"encoder\", encoder_name)\n",
    "        mlflow.log_param(\"arch\", arch)\n",
    "        mlflow.log_param(\"learning rate\", lr)\n",
    "        mlflow.set_tag(\"lakefs_demos\", \"image_segmentation\")\n",
    "        with converter_train.make_torch_dataloader(num_epochs=1, transform_spec=get_transform_spec(is_train=True),\n",
    "                                                   batch_size=BATCH_SIZE) as train_dataloader, \\\n",
    "                converter_test.make_torch_dataloader(num_epochs=1, transform_spec=get_transform_spec(is_train=False),\n",
    "                                                     batch_size=BATCH_SIZE) as valid_dataloader:\n",
    "            trainer.fit(model, train_dataloaders=train_dataloader, val_dataloaders=valid_dataloader)\n",
    "            delattr(model, \"trainer\")\n",
    "            # add model requirement\n",
    "            reqs = mlflow.pytorch.get_default_pip_requirements() + [\n",
    "                \"git+https://github.com/qubvel/segmentation_models.pytorch\", \"pytorch-lightning==\" + pl.__version__]\n",
    "            mlflow.pyfunc.log_model(artifact_path=\"model\", python_model=CVModelWrapper(model), pip_requirements=reqs)\n",
    "            # log and returns model accuracy\n",
    "            valid_metrics = trainer.validate(model, dataloaders=valid_dataloader, verbose=False)\n",
    "            valid_per_image_iou = valid_metrics[0]['valid_per_image_iou']\n",
    "            mlflow.log_metric(\"valid_per_image_iou\", valid_per_image_iou)\n",
    "            mlflow.log_metric(\"loss\", 1 - valid_per_image_iou)\n",
    "            \n",
    "            # Log information about dataset and lakeFS tags/commits\n",
    "            if localOrDistributedComputing == \"LOCAL\":\n",
    "                dataset = mlflow.data.load_delta(path=f\"{gold_data_path}/{training_data_folder}\")\n",
    "                mlflow.log_input(dataset, context=\"Gold Dataset\")\n",
    "            lakefs_dataset_tag = f\"{lakefsUIEndPoint}/repositories/{repo_name}/objects?ref={goldDatasetTagID}&path=gold/train_v2/\"\n",
    "            dictionary = {\"lakefs_dataset\": lakefs_dataset_tag}\n",
    "            mlflow.log_dict(dictionary, \"model/lakefs_dataset.json\")\n",
    "            mlflow.set_tag(\"lakefs_repo\", repo_name)\n",
    "            mlflow.set_tag(\"lakefs_branch\", experimentBranchN)\n",
    "            mlflow.set_tag(\"lakefs_dataset\", lakefs_dataset_tag)\n",
    "\n",
    "            return valid_per_image_iou"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "a6065c3e-f3a4-4fb1-ab11-4caa5289fa10",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Other helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "4d585cb4-0881-4fbe-9310-f589b6a96f29",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def commit_metadata_for_best_model(best_model, model_registered):\n",
    "    if localOrDistributedComputing == \"LOCAL\":\n",
    "        commitMetadata='::lakefs::Registered Model::url[url:ui]=http://127.0.0.1:5001/#/models/' + model_registered.name + '/versions/' + str(model_registered.version) + \\\n",
    "            ',Model Name=' + model_registered.name + \\\n",
    "            ',Model Version=' + str(model_registered.version) + \\\n",
    "            ',Model Run Name=' + best_model[\"tags.mlflow.runName\"] + \\\n",
    "            ',Model Metric valid_per_image_iou=' + str(best_model[\"metrics.valid_per_image_iou\"])\n",
    "    elif localOrDistributedComputing == \"DISTRIBUTED\":    \n",
    "        commitMetadata={'::lakefs::Registered Model::url[url:ui]' : 'https://'+spark.conf.get(\"spark.databricks.workspaceUrl\")+'/#mlflow/models/' + model_registered.name + '/versions/' + str(model_registered.version),\n",
    "                        'Model Name' : model_registered.name,\n",
    "                        'Model Version' : str(model_registered.version),\n",
    "                        'Model Run Name' : best_model[\"tags.mlflow.runName\"],\n",
    "                        'Model Metric valid_per_image_iou' : str(best_model[\"metrics.valid_per_image_iou\"])                \n",
    "                       }\n",
    "    return commitMetadata                                                                                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6d7fcd85-1afa-42f7-b3ff-127c564c952f",
     "showTitle": false,
     "title": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def display_gold_images_header():\n",
    "    wi1 = widgets.Label(value='image_id', disabled=True, layout=widgets.Layout(width='120px'))\n",
    "    wi2 = widgets.Label(value='boats', disabled=True, layout=widgets.Layout(width='60px'))\n",
    "    wi3 = widgets.Label(value='mask', disabled=True, layout=widgets.Layout(width='150px'))\n",
    "    wi4 = widgets.Label(value='content', disabled=True, layout=widgets.Layout(width='150px'))\n",
    "    wid=widgets.HBox([wi1,wi2,wi3,wi4])\n",
    "    display(wid)\n",
    "    \n",
    "def display_gold_images_row(row):\n",
    "    wi1 = widgets.Text(value=row['image_id'], disabled=True, layout=widgets.Layout(width='120px'))\n",
    "    wi2 = widgets.IntText(value=row['boat_number'], disabled=True, layout=widgets.Layout(width='60px'))\n",
    "    wi3 = widgets.Image(value=row['mask'], format='jpg', width=150, height=150)\n",
    "    wi4 = widgets.Image(value=row['content'], format='jpg', width=150, height=150)\n",
    "    wid=widgets.HBox([wi1,wi2,wi3,wi4])\n",
    "    display(wid)\n",
    "    \n",
    "def display_gold_images_local(df_gold_images):\n",
    "    display_gold_images_header()\n",
    "    dataCollect = df_gold_images.collect()\n",
    "    for row in dataCollect:\n",
    "        display_gold_images_row(row)\n",
    "        \n",
    "def display_gold_images(df_gold_images):\n",
    "    if localOrDistributedComputing == \"LOCAL\":\n",
    "        display_gold_images_local(df_gold_images)\n",
    "    elif localOrDistributedComputing == \"DISTRIBUTED\":\n",
    "        display(df_gold_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "bc46b71a-592c-4ef7-8b1f-225d43c23ced",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# This works with Jupyter notebook but does not work with Databricks notebook\n",
    "def display_capture(capture, lines):\n",
    "    capture_split = capture.stdout.split('\\n')\n",
    "    if lines > len(capture_split):\n",
    "        lines = len(capture_split)\n",
    "    for x in range(lines):\n",
    "        print(capture_split[-(lines-x)])\n",
    "        \n",
    "    capture_split = capture.stderr.split('\\n')\n",
    "    if lines > len(capture_split):\n",
    "        lines = len(capture_split)\n",
    "    for x in range(lines):\n",
    "        print(capture_split[-(lines-x)])"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "ImageSegmentationSetup",
   "widgets": {}
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}