notebooks/samples/ConditionalKNN - Exploring Art Across Cultures.ipynb

{
   "cells":[
      {
         "cell_type":"markdown",
         "source":[
            "# Exploring Art across Culture and Medium with Fast, Conditional, k-Nearest Neighbors\n\n<img src=\"https://mmlspark.blob.core.windows.net/graphics/art/cross_cultural_matches.jpg\"  width=\"600\"/>\n\nThis notebook serves as a guideline for match-finding via k-nearest-neighbors. In the code below, we will set up code that allows queries involving cultures and mediums of art amassed from the Metropolitan Museum of Art in NYC and the Rijksmuseum in Amsterdam."
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"markdown",
         "source":[
            "### Overview of the BallTree\nThe structure functioning behind the kNN model is a BallTree, which is a recursive binary tree where each node (or \"ball\") contains a partition of the points of data to be queried. Building a BallTree involves assigning data points to the \"ball\" whose center they are closest to (with respect to a certain specified feature), resulting in a structure that allows binary-tree-like traversal and lends itself to finding k-nearest neighbors at a BallTree leaf."
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"markdown",
         "source":[
            "#### Setup\nImport necessary Python libraries and prepare dataset."
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"code",
         "source":[
            "from pyspark.sql.types import *\nfrom pyspark.ml.feature import Normalizer\nfrom pyspark.sql.functions import lit, array, array_contains, udf, col, struct\nfrom mmlspark.nn import ConditionalKNN, ConditionalKNNModel\nfrom PIL import Image\nfrom io import BytesIO\n\nimport requests\nimport numpy as np\nimport matplotlib.pyplot as plt"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":4
      },
      {
         "cell_type":"markdown",
         "source":[
            "Our dataset comes from a table containing artwork information from both the Met and Rijks museums. The schema is as follows:\n\n- **id**: A unique identifier for a piece of art\n  - Sample Met id: *388395* \n  - Sample Rijks id: *SK-A-2344* \n- **Title**: Art piece title, as written in the museum's database\n- **Artist**: Art piece artist, as written in the museum's database\n- **Thumbnail_Url**: Location of a JPEG thumbnail of the art piece\n- **Image_Url** Location of an image of the art piece hosted on the Met/Rijks website\n- **Culture**: Category of culture that the art piece falls under\n  - Sample culture categories: *latin american*, *egyptian*, etc\n- **Classification**: Category of medium that the art piece falls under\n  - Sample medium categories: *woodwork*, *paintings*, etc\n- **Museum_Page**: Link to the work of art on the Met/Rijks website\n- **Norm_Features**: Embedding of the art piece image\n- **Museum**: Specifies which museum the piece originated from"
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"code",
         "source":[
            "# loads the dataset and the two trained CKNN models for querying by medium and culture\ndf = spark.read.parquet(\"wasbs://publicwasb@mmlspark.blob.core.windows.net/met_and_rijks.parquet\")\ndisplay(df.drop(\"Norm_Features\"))"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":6
      },
      {
         "cell_type":"markdown",
         "source":[
            "#### Define categories to be queried on\nWe will be using two kNN models: one for culture, and one for medium. The categories for each grouping are defined below."
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"code",
         "source":[
            "from pyspark.sql.types import BooleanType\n\n#mediums = ['prints', 'drawings', 'ceramics', 'textiles', 'paintings', \"musical instruments\",\"glass\", 'accessories', 'photographs',  \"metalwork\", \n#           \"sculptures\", \"weapons\", \"stone\", \"precious\", \"paper\", \"woodwork\", \"leatherwork\", \"uncategorized\"]\n\nmediums = ['paintings', 'glass', 'ceramics']\n\n#cultures = ['african (general)', 'american', 'ancient american', 'ancient asian', 'ancient european', 'ancient middle-eastern', 'asian (general)', \n#            'austrian', 'belgian', 'british', 'chinese', 'czech', 'dutch', 'egyptian']#, 'european (general)', 'french', 'german', 'greek', \n#            'iranian', 'italian', 'japanese', 'latin american', 'middle eastern', 'roman', 'russian', 'south asian', 'southeast asian', \n#            'spanish', 'swiss', 'various']\n\ncultures = ['japanese', 'american', 'african (general)']\n\n# Uncomment the above for more robust and large scale searches!\n\nclasses = cultures + mediums\n\nmedium_set = set(mediums)\nculture_set = set(cultures)\nselected_ids = {\"AK-RBK-17525-2\", \"AK-MAK-1204\", \"AK-RAK-2015-2-9\"}\n\nsmall_df = df.where(udf(lambda medium, culture, id_val: (medium in medium_set) or (culture in culture_set) or (id_val in selected_ids), BooleanType())(\"Classification\", \"Culture\", \"id\"))\n\nsmall_df.count()"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":8
      },
      {
         "cell_type":"markdown",
         "source":[
            "### Define and fit ConditionalKNN models\nBelow, we create ConditionalKNN models for both the medium and culture columns; each model takes in an output column, features column (feature vector), values column (cell values under the output column), and label column (the quality that the respective KNN is conditioned on)."
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"code",
         "source":[
            "medium_cknn = (ConditionalKNN()\n  .setOutputCol(\"Matches\")\n  .setFeaturesCol(\"Norm_Features\")\n  .setValuesCol(\"Thumbnail_Url\")\n  .setLabelCol(\"Classification\")\n  .fit(small_df))"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":10
      },
      {
         "cell_type":"code",
         "source":[
            "culture_cknn = (ConditionalKNN()\n  .setOutputCol(\"Matches\")\n  .setFeaturesCol(\"Norm_Features\")\n  .setValuesCol(\"Thumbnail_Url\")\n  .setLabelCol(\"Culture\")\n  .fit(small_df))"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":11
      },
      {
         "cell_type":"markdown",
         "source":[
            "#### Define matching and visualizing methods\n\nAfter the intial dataset and category setup, we prepare methods that will query and visualize the conditional kNN's results. \n\n`addMatches()` will create a Dataframe with a handful of matches per category."
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"code",
         "source":[
            "def add_matches(classes, cknn, df):\n  results = df\n  for label in classes:\n    results = (cknn.transform(results.withColumn(\"conditioner\", array(lit(label))))\n                 .withColumnRenamed(\"Matches\", \"Matches_{}\".format(label)))\n  return results"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":13
      },
      {
         "cell_type":"markdown",
         "source":[
            "`plot_urls()` calls `plot_img` to visualize top matches for each category into a grid."
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"code",
         "source":[
            "def plot_img(axis, url, title):\n  response = requests.get(url)\n  img = Image.open(BytesIO(response.content)).convert('RGB')\n  axis.imshow(img, aspect=\"equal\")\n  if title is not None: axis.set_title(title,  fontsize=4)\n  axis.axis(\"off\")\n\ndef plot_urls(url_arr, titles, filename):\n  nx, ny = url_arr.shape\n  \n  plt.figure(figsize=(nx*5, ny*5), dpi=1600)\n  fig, axes = plt.subplots(ny,nx)\n  \n  # reshape required in the case of 1 image query\n  if len(axes.shape) == 1:\n    axes = axes.reshape(1, -1)\n    \n  for i in range(nx):\n     for j in range(ny):\n          if j == 0:\n            plot_img(axes[j, i], url_arr[i,j], titles[i])\n          else:\n            plot_img(axes[j, i],  url_arr[i,j], None)\n            \n  plt.savefig(filename, dpi=1600) # saves the results as a PNG\n\n  display(plt.show())"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":15
      },
      {
         "cell_type":"markdown",
         "source":[
            "### Putting it all together\nBelow, we define `test_all()` to take in the data, CKNN models, the art id values to query on, and the file path to save the output visualization to. The medium and culture models were previously trained and loaded."
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"code",
         "source":[
            "# main method to test a particular dataset with two CKNN models and a set of art IDs, saving the result to filename.png\n\ndef test_all(data, cknn_medium, cknn_culture, test_ids, root):\n  is_nice_obj = udf(lambda obj: obj in test_ids, BooleanType())\n  test_df = data.where(is_nice_obj(\"id\"))\n  \n  results_df_medium = add_matches(mediums, cknn_medium, test_df)\n  results_df_culture = add_matches(cultures, cknn_culture, results_df_medium)\n  \n  results = results_df_culture.collect()\n  \n  original_urls = [row[\"Thumbnail_Url\"] for row in results]\n  \n  culture_urls = [ [row[\"Matches_{}\".format(label)][0][\"value\"] for row in results] for label in cultures]\n  culture_url_arr = np.array([original_urls] + culture_urls)[:, :]\n  plot_urls(culture_url_arr, [\"Original\"] + cultures, root + \"matches_by_culture.png\")\n  \n  medium_urls = [ [row[\"Matches_{}\".format(label)][0][\"value\"] for row in results] for label in mediums]\n  medium_url_arr = np.array([original_urls] + medium_urls)[:, :]\n  plot_urls(medium_url_arr, [\"Original\"] + mediums, root + \"matches_by_medium.png\")\n  \n  return results_df_culture"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":17
      },
      {
         "cell_type":"markdown",
         "source":[
            "### Demo\nThe following cell performs batched queries given desired image IDs and a filename to save the visualization.\n\n\n<img src=\"https://mmlspark.blob.core.windows.net/graphics/art/cross_cultural_matches.jpg\"  width=\"600\"/>"
         ],
         "metadata":{

         }
      },
      {
         "cell_type":"code",
         "source":[
            "# sample query\nresult_df = test_all(small_df, medium_cknn, culture_cknn, selected_ids, root=\".\")"
         ],
         "metadata":{

         },
         "outputs":[

         ],
         "execution_count":19
      }
   ],
    "metadata": {
     "anaconda-cloud": {},
     "kernelspec": {
      "display_name": "Python [default]",
      "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.6.3"
     }
    },
    "nbformat": 4,
    "nbformat_minor": 2

}