notebooks/process_chronam_pages.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This Newspaper Navigator Dataset Pipeline\n",
    "\n",
    "By Benjamin Charles Germain Lee (2020 Library of Congress Innovator-in-Residence)\n",
    "\n",
    "Using the manifests saved in the repo, this notebook iterates over each manifest (corresponding to a Chronicling America batch and performs the following.  *Note that if you would like to use an updated version of the manifests, please see the repo \"chronam-get-images\" (https://github.com/bcglee/chronam-get-images) for how to generate the manifests.*\n",
    "\n",
    "1. Systematically downloads the images and METS/ALTO OCR (XML) from the proper S3 buckets. This step allows for CPU multiprocessing; adjust N_CPU_PROCESSES as desired in main(). *Note: if you're experimenting with this pipeline and can't configure access to the Chronicling America S3 buckets, you can use the repo chronam-get-images (https://github.com/bcglee/chronam-get-images) to pull down the JPG and XML files of desired pages; then, move those files to '../chronam_files').*\n",
    "2. Generates and saves predictions for each JPG image. This step utilizes model weights generated using the notebook 'train_model.ipynb'. This step uses all available GPUs.\n",
    "3. Adds captions from the METS/ALTO OCR in the XML for each image as metadata. This step is performed by identifying text within each predicted bounding box. This step allows for CPU multiprocessing.\n",
    "4. Crops all of the predicted visual content and saves the cropped images. This step allows for CPU multiprocessing.\n",
    "5. Generates embeddings for the predicted visual content for each image and adds the embeddings to the metadata. Currently, img2vec is being utilized for this (https://github.com/christiansafka/img2vec). This step uses all available GPUs.\n",
    "6. Packages the files according to the Chronicling America file structure and sends the files to an S3 bucket. The downloaded files and generated metadata are all then deleted to free up space.\n",
    "\n",
    "The pipeline is driven by main(), which is the last cell in the notebook. The functions used in main() appear below with short descriptions. Though this workflow differs from most Jupyter notebooks, the notebook format is being used here for the ease of annotating code.\n",
    "\n",
    "NOTE: if you would like to run this code, you MUST save the notebook as a Python file using the command \"jupyter nbconvert --to script process_chronam_pages.ipynb\" and run the Python script using \"python process_chronam_pages.py\".  This is necessary because the notebook is unable to handle multiprocessing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The next two cells include the code for the systematic download of Chroncling America images from the S3 buckets.\n",
    "\n",
    "This first cell handles imports and initial settings for pulling down the files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import botocore\n",
    "import s3fs\n",
    "import glob\n",
    "import sys\n",
    "import os\n",
    "import time\n",
    "from PIL import Image\n",
    "import io\n",
    "import math\n",
    "import datetime\n",
    "import subprocess\n",
    "import collections\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This second cell handles the function for file retrieval for a specified manifest and destination directory.  \n",
    "\n",
    "Note that we test first to see if a JPG exists in 'ndnp-jpeg-surrogates' and if not, we then grab the JP2 from 'ndnp-batches' (converting the JP2 to JPG requires overhead). The XML file is then downloaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that retrieves .jpg and .xml files for each filepath in manifest\n",
    "def retrieve_files(packet):\n",
    "    \n",
    "    # sets boto3 to run with s3\n",
    "    s3 = boto3.resource('s3')\n",
    "    \n",
    "    # use s3fs file system for checking file existence\n",
    "    s3fs_filye_sys = s3fs.S3FileSystem()        \n",
    "\n",
    "    # creates dict for storing widths/heights of images\n",
    "    im_size_dict = {}  \n",
    "    \n",
    "    # grab directory to CD into first (it is the firs entry in the array)\n",
    "    dir_path = packet[0]\n",
    "    os.chdir(dir_path)\n",
    "    \n",
    "    # grabs page_filepaths from the data packet\n",
    "    page_filepaths = packet[1]\n",
    "        \n",
    "    # iterate through each filepath and download\n",
    "    for page_filepath in page_filepaths:\n",
    "                \n",
    "        # sets filepath for download destination (note: file is .jp2, so we need to replace suffixes below)\n",
    "        local_filepath = page_filepath.replace('/', '_')\n",
    "\n",
    "        # see: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/migrations3.html\n",
    "        # see also:  https://www.edureka.co/community/17558/python-aws-boto3-how-do-i-read-files-from-s3-bucket\n",
    "        try:\n",
    "\n",
    "            # if the JPEG exists in ndnp-jpeg-surrogates, pull down from that S3 bucket\n",
    "            if s3fs_filye_sys.exists('ndnp-jpeg-surrogates/' + page_filepath.replace(\".jp2\", \".jpg\")):\n",
    "                obj = s3.Object('ndnp-jpeg-surrogates', page_filepath.replace(\".jp2\", \".jpg\"))\n",
    "                body = obj.get()['Body'].read()\n",
    "                im = Image.open(io.BytesIO(body))\n",
    "                im.resize((math.floor(im.width/6), math.floor(im.height/6)), resample=0).save(local_filepath.replace(\".jp2\", \".jpg\"))\n",
    "                im_size_dict[local_filepath.replace(\".jp2\", \".jpg\")] = (im.width, im.height)\n",
    "\n",
    "            # if the JPEG doesn't exist, pull down the JPEG-2000 from 'ndnp-batches' S3 bucket and convert to JPG\n",
    "            else:\n",
    "                s3.Bucket('ndnp-batches').download_file(page_filepath, local_filepath)\n",
    "                subprocess.call('gm convert ' + local_filepath + ' ' + local_filepath.replace(\".jp2\", \".jpg\"), shell=True)\n",
    "                subprocess.call('rm ' + local_filepath, shell=True)\n",
    "\n",
    "                if not os.path.exists(local_filepath.replace(\".jp2\", \".jpg\")):\n",
    "                    continue\n",
    "\n",
    "                im = Image.open(local_filepath.replace(\".jp2\", \".jpg\"))\n",
    "                im.resize((math.floor(im.width/6), math.floor(im.height/6)), resample=0).save(local_filepath.replace(\".jp2\", \".jpg\"))\n",
    "                im_size_dict[local_filepath.replace(\".jp2\", \".jpg\")] = (im.width, im.height)\n",
    "\n",
    "        except botocore.exceptions.ClientError as e:\n",
    "            if e.response['Error']['Code'] == \"404\":\n",
    "                print(\"The object does not exist.\")\n",
    "            else:\n",
    "                print(\"Error in downloading JP2.\")\n",
    "                raise\n",
    "                \n",
    "        try:\n",
    "            s3.Bucket('ndnp-batches').download_file(page_filepath.replace(\".jp2\", \".xml\"), local_filepath.replace(\".jp2\", \".xml\"))\n",
    "        except botocore.exceptions.ClientError as e:\n",
    "            if e.response['Error']['Code'] == \"404\":\n",
    "                print(\"The object does not exist.\")\n",
    "            else:\n",
    "                print(\"Error in downloading XML.\")\n",
    "                raise\n",
    "        \n",
    "    return im_size_dict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The next two cells load the finetuned model and define the function for performing predictions on the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import some common libraries\n",
    "import cv2\n",
    "import random\n",
    "import json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "# import deep learning imports\n",
    "import detectron2\n",
    "import torch\n",
    "from detectron2.config import get_cfg\n",
    "from detectron2.data import DatasetCatalog, MetadataCatalog\n",
    "from detectron2.data.datasets import register_coco_instances\n",
    "from detectron2.engine import DefaultTrainer\n",
    "from detectron2.engine import DefaultPredictor\n",
    "from detectron2.evaluation import COCOEvaluator\n",
    "from detectron2.utils.visualizer import Visualizer\n",
    "from detectron2.utils.visualizer import ColorMode\n",
    "from detectron2.utils.logger import setup_logger\n",
    "from detectron2.modeling import build_model\n",
    "from detectron2.checkpoint import DetectionCheckpointer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_predictions(zipped):\n",
    "            \n",
    "    # unzips packed information for process to perform predictions\n",
    "    \n",
    "    S3_SAVE_DIR = zipped[0]\n",
    "    OUTPUT_SAVE_DIR = zipped[1]\n",
    "    dir_name = zipped[2]\n",
    "    INFERENCE_BATCH_SIZE = zipped[3]\n",
    "    filepaths = zipped[4]\n",
    "    ID = zipped[5]\n",
    "\n",
    "    with torch.cuda.device(ID):\n",
    "\n",
    "        # navigates to correct directory (process is spawned in /notebooks)\n",
    "\n",
    "        os.chdir(S3_SAVE_DIR + dir_name)\n",
    "\n",
    "        # sets up model for process\n",
    "\n",
    "        setup_logger()\n",
    "        cfg = get_cfg()\n",
    "        cfg.merge_from_file(\"../../..//detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml\")\n",
    "\n",
    "        # sets prediction score threshold - this is commented out and defaults to 0.05 in Detectron2\n",
    "        # if you would like to adjust the threshold, uncomment and set to the desired value in [0, 1)\n",
    "        # cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5\n",
    "\n",
    "        # sets number of object classes to 7\n",
    "        # (\"Illustration/Photograph\", \"Photograph\", \"Comics/Cartoon\", \"Editorial Cartoon\", \"Map\", \"Headline\", \"Ad\")\n",
    "        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 7\n",
    "\n",
    "        # build model\n",
    "        model = build_model(cfg)\n",
    "\n",
    "        # see:  https://github.com/facebookresearch/detectron2/issues/282 \n",
    "        # (must load weights this way if using model)\n",
    "        DetectionCheckpointer(model).load(\"../../model_weights/model_final.pth\")\n",
    "        model.train(False) \n",
    "\n",
    "        # construct batches\n",
    "        batches = chunk(filepaths, math.ceil(len(filepaths)/INFERENCE_BATCH_SIZE))\n",
    "\n",
    "        # iterate through images\n",
    "        for batch in batches:\n",
    "\n",
    "            # sets up inputs by loading in all files in batch\n",
    "            inputs = []\n",
    "\n",
    "            # stores image dimensions\n",
    "            dimensions = []\n",
    "\n",
    "            # iterate through files in batch\n",
    "            for file in batch:\n",
    "\n",
    "                # read in image\n",
    "                image = cv2.imread(file)\n",
    "\n",
    "                # store image dimensions\n",
    "                height, width, _ = image.shape\n",
    "                dimensions.append([width, height])\n",
    "\n",
    "                # perform inference on batch\n",
    "                image = np.transpose(image,(2,0,1))\n",
    "                # see https://github.com/facebookresearch/detectron2/issues/282 for in-depth description of why \n",
    "                # image is loaded in this way\n",
    "                image_tensor = torch.from_numpy(image)\n",
    "                inputs.append({\"image\": image_tensor})\n",
    "\n",
    "            # performs inference\n",
    "            outputs = model(inputs)\n",
    "\n",
    "            # saves predictions\n",
    "            predictions = {}\n",
    "\n",
    "            # iterate over images in batch and save predictions to JSON\n",
    "            for i in range(0, len(batch)):\n",
    "\n",
    "                # saves filepath in format of ChronAm file structure\n",
    "                predictions[\"filepath\"] = dir_name + \"data/\" + batch[i].split(\"data_\")[1].replace(dir_name, '').replace('_', '/').replace('.jpg', '.jp2')\n",
    "\n",
    "                # parses metadata from filepath\n",
    "                date_str = predictions[\"filepath\"].split('/')[-2]\n",
    "                predictions[\"batch\"] = dir_name[:-1]\n",
    "                predictions[\"lccn\"] = predictions[\"filepath\"].split('/')[-4]\n",
    "                predictions[\"pub_date\"] = str(datetime.date(int(date_str[:4]), int(date_str[4:6]), int(date_str[6:8])))\n",
    "                predictions[\"edition_seq_num\"] = int(date_str[8:10])\n",
    "                if \"a\" in predictions[\"filepath\"].split('/')[-1][:-4] or \"b\" in predictions[\"filepath\"].split('/')[-1][:-4]:\n",
    "                    continue\n",
    "                predictions[\"page_seq_num\"] = int(predictions[\"filepath\"].split('/')[-1][:-4])\n",
    "                \n",
    "                # saves predictions\n",
    "                # we first normalize the bounding box coordinates\n",
    "                boxes = outputs[i][\"instances\"].get_fields()[\"pred_boxes\"].to(\"cpu\").tensor.tolist()\n",
    "                normalized_boxes = []\n",
    "                width = dimensions[i][0]\n",
    "                height = dimensions[i][1]\n",
    "\n",
    "                for box in boxes:\n",
    "                    normalized_box = (box[0]/float(width), box[1]/float(height), box[2]/float(width), box[3]/float(height))\n",
    "                    normalized_boxes.append(normalized_box)\n",
    "\n",
    "                # saves additional outputs of predictions\n",
    "                predictions[\"boxes\"] = normalized_boxes\n",
    "                predictions[\"scores\"] = outputs[i][\"instances\"].get_fields()[\"scores\"].to(\"cpu\").tolist()\n",
    "                predictions[\"pred_classes\"] = outputs[i][\"instances\"].get_fields()[\"pred_classes\"].to(\"cpu\").tolist()\n",
    "\n",
    "                with open(OUTPUT_SAVE_DIR + dir_name + batch[i].replace('.jpg','.json'), \"w\") as fp:\n",
    "                    json.dump(predictions, fp)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The next two cells define functions for extracting the OCR within each predicted box.\n",
    "\n",
    "1. The first cell defines the function for returning the proper OCR for a specific page.\n",
    "2. The second cell defines the function for iterating over the JSON files containing the predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import xml.etree.cElementTree as ET\n",
    "# etree\n",
    "from lxml import etree as ET\n",
    "from xml.etree.ElementTree import ElementTree\n",
    "\n",
    "# tolerance around box for testing whether OCR falls within bounds\n",
    "WIDTH_TOLERANCE = 0.000\n",
    "HEIGHT_TOLERANCE = 0.000\n",
    "\n",
    "# given a file path and a list of bounding boxes, this function traverses the associated XML\n",
    "# and returns the OCR within each bounding box\n",
    "def retrieve_ocr_for_file(xml_filepath, true_img_filepath, page_width_pix, page_height_pix, bounding_boxes, predicted_classes):\n",
    "\n",
    "    # creates empty nested list fo storing OCR in each box\n",
    "    ocr = [ [] for i in range(len(bounding_boxes)) ]\n",
    "\n",
    "    # sets tree and root based on filepath\n",
    "    parser = ET.XMLParser()\n",
    "    tree = ET.parse(xml_filepath, parser)\n",
    "    root = tree.getroot()\n",
    "    \n",
    "    # sets tag prefix (everywhere)\n",
    "    prefix = root.tag.split('}')[0] + '}'\n",
    "\n",
    "    # traverses to layout and then the page and then the print space\n",
    "    layout = root.find(prefix + 'Layout')\n",
    "    page = layout.find(prefix + 'Page')\n",
    "    print_space = page.find(prefix + 'PrintSpace')\n",
    "    \n",
    "    if print_space is None:\n",
    "        return ocr\n",
    "\n",
    "    text_boxes =  [textblock for textblock in print_space.iterchildren(prefix + \"TextBlock\")]\n",
    "    \n",
    "    # gets page height and page width in inch1200 units\n",
    "    page_width_inch = int(page.attrib['WIDTH'])\n",
    "    page_height_inch = int(page.attrib['HEIGHT'])\n",
    "\n",
    "    # sets conversion to normalized coordinates for comparison between METS/ALTO and predicted boxes\n",
    "    W_CONVERSION = 1./float(page_width_inch)\n",
    "    H_CONVERSION = 1./float(page_height_inch)\n",
    "\n",
    "    if page_width_inch == 0 or page_height_inch == 0:\n",
    "        return ocr\n",
    "\n",
    "    # we now iterate over each bounding box\n",
    "    for i in range(0, len(bounding_boxes)):\n",
    "\n",
    "        bounding_box = bounding_boxes[i]\n",
    "        predicted_class = predicted_classes[i]\n",
    "\n",
    "        # we then iterate over each text box\n",
    "        for text_box in text_boxes:\n",
    "                        \n",
    "            box_w1 = int(float(text_box.attrib[\"HPOS\"]))\n",
    "            box_h1 = int(float(text_box.attrib[\"VPOS\"]))\n",
    "            box_w2 = box_w1 + int(float(text_box.attrib[\"WIDTH\"]))\n",
    "            box_h2 = box_h1 + int(float(text_box.attrib[\"HEIGHT\"]))\n",
    "            \n",
    "            # if the text box and bounding box do not intersect, we skip (as no text will overlap in smaller units)\n",
    "            if box_w2*W_CONVERSION < bounding_box[0] and box_h2*H_CONVERSION < bounding_box[1]:\n",
    "                continue\n",
    "            if box_w1*W_CONVERSION > bounding_box[0] + bounding_box[2] and box_h2*H_CONVERSION < bounding_box[1]:\n",
    "                continue\n",
    "            if box_w2*W_CONVERSION < bounding_box[0] and box_h1*H_CONVERSION > bounding_box[1] + bounding_box[3]:\n",
    "                continue\n",
    "            if box_w1*W_CONVERSION > bounding_box[0] + bounding_box[2] and box_h1*H_CONVERSION > bounding_box[1] + bounding_box[3]:\n",
    "                continue\n",
    "                \n",
    "            # we then iterate over the text lines in each box\n",
    "            for text_line in text_box.iterchildren(prefix + 'TextLine'):\n",
    "                \n",
    "                line_w1 = int(float(text_box.attrib[\"HPOS\"]))\n",
    "                line_h1 = int(float(text_box.attrib[\"VPOS\"]))\n",
    "                line_w2 = line_w1 + int(float(text_box.attrib[\"WIDTH\"]))\n",
    "                line_h2 = line_h1 + int(float(text_box.attrib[\"HEIGHT\"]))\n",
    "\n",
    "                # if the text box and bounding box do not intersect, we skip (as no text will overlap in smaller units)\n",
    "                if line_w2*W_CONVERSION < bounding_box[0] and line_h2*H_CONVERSION < bounding_box[1]:\n",
    "                    continue\n",
    "                if line_w1*W_CONVERSION > bounding_box[0] + bounding_box[2] and line_h2*H_CONVERSION < bounding_box[1]:\n",
    "                    continue\n",
    "                if line_w2*W_CONVERSION < bounding_box[0] and line_h1*H_CONVERSION > bounding_box[1] + bounding_box[3]:\n",
    "                    continue\n",
    "                if line_w1*W_CONVERSION > bounding_box[0] + bounding_box[2] and line_h1*H_CONVERSION > bounding_box[1] + bounding_box[3]:\n",
    "                    continue\n",
    "                \n",
    "                # we now iterate over every string in each line (each string is separated by whitespace)\n",
    "                for string in text_line.iterchildren(prefix + 'String'):\n",
    "            \n",
    "                    w1 = int(float(string.attrib[\"HPOS\"]))\n",
    "                    h1 = int(float(string.attrib[\"VPOS\"]))\n",
    "                    w2 = w1 + int(float(string.attrib[\"WIDTH\"]))\n",
    "                    h2 = h1 + int(float(string.attrib[\"HEIGHT\"]))\n",
    "\n",
    "                    # checks if the text appears within the bounding box & extra tolerance for words that are clipped\n",
    "                    if w1*W_CONVERSION > bounding_box[0] - WIDTH_TOLERANCE:\n",
    "                        if w2*W_CONVERSION < bounding_box[2] + WIDTH_TOLERANCE:\n",
    "                            if h1*H_CONVERSION > bounding_box[1] - HEIGHT_TOLERANCE:\n",
    "                                if h2*H_CONVERSION < bounding_box[3] + HEIGHT_TOLERANCE:\n",
    "\n",
    "                                    # appends text content to list\n",
    "                                    ocr[i].append(string.attrib[\"CONTENT\"])\n",
    "\n",
    "    return ocr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieve_ocr(packet):\n",
    "\n",
    "    # grab contents of packet, CD into correct directory\n",
    "    dir_name = packet[1]\n",
    "    os.chdir(packet[0] + dir_name)\n",
    "    json_info = packet[2]\n",
    "\n",
    "    # we now iterate through all of the predictions JSON files\n",
    "    for json_entry in json_info:\n",
    "        \n",
    "        # unpacks the input from Pool\n",
    "        json_filepath = json_entry[0]\n",
    "        im_width = json_entry[1]\n",
    "        im_height = json_entry[2]\n",
    "        \n",
    "        # loads the JSON\n",
    "        with open(json_filepath) as f:\n",
    "            predictions = json.load(f)\n",
    "        \n",
    "        # pulls off relevant data fields from the JSON\n",
    "        original_img_filepath = predictions['filepath']\n",
    "        boxes = predictions['boxes']\n",
    "        scores = predictions['scores']\n",
    "        classes = predictions['pred_classes']\n",
    "\n",
    "        # sets the number of predicted bounding boxes\n",
    "        n_pred = len(scores)\n",
    "\n",
    "        # we now find the XML and JPG files corresponding to this predictions JSON\n",
    "        xml_filepath = S3_SAVE_DIR + dir_name + json_filepath.replace('.json', '.xml')\n",
    "        jpg_filepath = S3_SAVE_DIR + dir_name + json_filepath.replace('.json', '.jpg')\n",
    "\n",
    "        # stores list of OCR\n",
    "        ocr = []\n",
    "\n",
    "        # we only try to retrieve the OCR if there is one or more predicted box\n",
    "        if n_pred > 0:\n",
    "            ocr = retrieve_ocr_for_file(xml_filepath, jpg_filepath, im_width, im_height, boxes, classes)\n",
    "\n",
    "        # adds the ocr field to the JSON metadata for the page\n",
    "        predictions['ocr'] = ocr\n",
    "\n",
    "        # we save the updated JSON\n",
    "        with open(json_filepath, 'w') as f:\n",
    "            json.dump(predictions, f)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This cell defines a function for cropping all of the predicted visual content:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop(packet):\n",
    "            \n",
    "    OUTPUT_SAVE_DIR = packet[0]\n",
    "    S3_SAVE_DIR = packet[1]\n",
    "    dir_name = packet[2]\n",
    "    json_filepaths = packet[3]\n",
    "    \n",
    "    os.chdir(OUTPUT_SAVE_DIR+dir_name)\n",
    "\n",
    "    for json_filepath in json_filepaths:\n",
    "        \n",
    "        # we load the JSON\n",
    "        with open(json_filepath) as f:\n",
    "            predictions = json.load(f)\n",
    "          \n",
    "        # load in boxes\n",
    "        boxes = predictions['boxes']\n",
    "        scores = predictions['scores']\n",
    "        classes = predictions['pred_classes']\n",
    "        \n",
    "        # grab filepath of image\n",
    "        jpg_filepath = S3_SAVE_DIR + dir_name + json_filepath.replace('.json', '.jpg')\n",
    "\n",
    "        # open image\n",
    "        im = Image.open(jpg_filepath)\n",
    "        \n",
    "        # empty list for storing embeddings\n",
    "        img_embeddings = []\n",
    "        \n",
    "        # empty list or storing filepaths of extracted visual content\n",
    "        content_filepaths = []\n",
    "\n",
    "        # iterate through boxes, crop, and send to embedding\n",
    "        for i in range(0, len(boxes)):\n",
    "            box = boxes[i]\n",
    "            pred_class = classes[i]\n",
    "            score = scores[i]\n",
    "            \n",
    "            # if it's a headline or the confidence score is less than 0.5, we skip the cropping\n",
    "            if pred_class == 5:\n",
    "                img_embeddings.append([])\n",
    "                content_filepaths.append([])\n",
    "                continue\n",
    "                \n",
    "            # crop image according to box (converted from normalized coordinates to image coordinates)\n",
    "            cropped = im.crop((box[0]*im.width, box[1]*im.height, box[2]*im.width, box[3]*im.height)).convert('RGB')\n",
    "            # save cropped image to output directory\n",
    "            cropped_filepath = json_filepath.replace(\".json\", \"_\" + str(i).zfill(3) + \"_\" + str(pred_class) + \"_\" + str(int(math.floor(100*score))).zfill(2) + \".jpg\")\n",
    "            cropped.save(cropped_filepath)\n",
    "            new_filepath = dir_name + \"data/\" + cropped_filepath.split(\"data_\")[1].replace(dir_name, '').replace('_', '/') \n",
    "            new_filepath = new_filepath[:new_filepath.rfind(\"/\")] + \"_\" + new_filepath[new_filepath.rfind(\"/\")+1:]\n",
    "            new_filepath = new_filepath[:new_filepath.rfind(\"/\")] + \"_\" + new_filepath[new_filepath.rfind(\"/\")+1:]\n",
    "            content_filepaths.append(new_filepath)\n",
    "\n",
    "        # add filepaths of extracted visual content to output\n",
    "        predictions['visual_content_filepaths'] = content_filepaths\n",
    "        \n",
    "        # we save the updated JSON\n",
    "        with open(json_filepath, 'w') as f:\n",
    "            json.dump(predictions, f)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This cell defines a function for generating embeddings of each predicted box:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from img2vec_pytorch import Img2Vec\n",
    "\n",
    "def generate_embeddings(zipped):\n",
    "    \n",
    "    # unzips packed information for process to perform predictions\n",
    "    \n",
    "    OUTPUT_SAVE_DIR = zipped[0]\n",
    "    S3_SAVE_DIR = zipped[1]\n",
    "    dir_name = zipped[2]\n",
    "    json_filepaths = zipped[3]\n",
    "    ID = zipped[4]\n",
    "\n",
    "    with torch.cuda.device(ID):\n",
    "\n",
    "        # load in img2vec\n",
    "        # we choose resnet embeddings\n",
    "        img2vec_resnet_50 = Img2Vec(cuda=True, model='resnet-50') \n",
    "        img2vec_resnet_18 = Img2Vec(cuda=True, model='resnet-18') \n",
    "    \n",
    "        # iterate through the JSON files\n",
    "        for json_filepath in json_filepaths:\n",
    "            \n",
    "            # we load the JSON\n",
    "            with open(json_filepath) as f:\n",
    "                predictions = json.load(f)\n",
    "\n",
    "            # load in boxes\n",
    "            boxes = predictions['boxes']\n",
    "            scores = predictions['scores']\n",
    "            classes = predictions['pred_classes']\n",
    "            cropped_filepaths = predictions['visual_content_filepaths']\n",
    "\n",
    "            # grab filepath of image\n",
    "            jpg_filepath = S3_SAVE_DIR + dir_name + json_filepath.replace('.json', '.jpg')\n",
    "\n",
    "            # empty list for storing embeddings\n",
    "            resnet_50_embeddings = []\n",
    "            resnet_18_embeddings = []\n",
    "\n",
    "            # iterate through boxes, crop, and send to embedding\n",
    "            for i in range(0, len(boxes)):\n",
    "\n",
    "                box = boxes[i]\n",
    "                pred_class = classes[i]\n",
    "                score = scores[i]\n",
    "                \n",
    "                # if it's a headline or confidence score is less than 0.5, we skip the embedding generation\n",
    "                if pred_class == 5 or score < 0.5:\n",
    "                    resnet_50_embeddings.append([])\n",
    "                    resnet_18_embeddings.append([])\n",
    "                    continue\n",
    "\n",
    "                cropped_filepath = cropped_filepaths[i]\n",
    "                # reformat to use flat file directory\n",
    "                cropped_filepath = cropped_filepath.replace(\"/\", \"_\")\n",
    "                \n",
    "                # open cropped image\n",
    "                im = Image.open(cropped_filepath).convert('RGB')\n",
    "                # generate embedding using img2vec\n",
    "                embedding_resnet_50 = img2vec_resnet_50.get_vec(im, tensor=False)\n",
    "                embedding_resnet_18 = img2vec_resnet_18.get_vec(im, tensor=False)\n",
    "                # add to list (render embedding numpy array as list to enable JSON serialization)\n",
    "                resnet_50_embeddings.append(embedding_resnet_50.tolist())\n",
    "                resnet_18_embeddings.append(embedding_resnet_18.tolist())\n",
    "                \n",
    "            embeddings_json = {}\n",
    "            embeddings_json['filepath'] = predictions['filepath']\n",
    "            embeddings_json['visual_content_filepaths'] = predictions['visual_content_filepaths']\n",
    "            # add embeddings to output\n",
    "            embeddings_json['resnet_50_embeddings'] = resnet_50_embeddings\n",
    "            embeddings_json['resnet_18_embeddings'] = resnet_18_embeddings\n",
    "\n",
    "            # we save the updated JSON\n",
    "            with open(json_filepath[:-5] + \"_embeddings.json\", 'w') as f:\n",
    "                json.dump(embeddings_json, f)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The cell below defines a function for chunking data for multiprocessing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that splits a list into n chunks for multiprocessing\n",
    "def chunk(file_list, n_chunks):\n",
    "    \n",
    "    # make chunks of files to be distributed across processes\n",
    "    chunks = []\n",
    "    chunk_size = math.ceil(float(len(file_list))/n_chunks)\n",
    "    for i in range(0, n_chunks-1):\n",
    "        chunks.append(file_list[i*chunk_size:(i+1)*chunk_size])\n",
    "    chunks.append(file_list[(n_chunks-1)*chunk_size:])\n",
    "    \n",
    "    return chunks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The cell below defines a function for validating that JP2 and XML filepaths derived from the manifests exist:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that determines whether the JP2 and XML files exist for specified files\n",
    "def files_exist(filepaths):\n",
    "    \n",
    "    no_jp2 = []\n",
    "    no_xml = []\n",
    "    good_filepaths = []\n",
    "    s3 = s3fs.S3FileSystem()\n",
    "    for filepath in filepaths:\n",
    "        if not s3.exists('ndnp-batches/' + filepath):\n",
    "            no_jp2.append(filepath)\n",
    "            continue\n",
    "        if not s3.exists('ndnp-batches/' + filepath.replace(\".jp2\", \".xml\")):\n",
    "            no_xml.append(filepath.replace(\".jp2\", \".xml\"))\n",
    "            continue\n",
    "        good_filepaths.append(filepath)\n",
    "\n",
    "    return [good_filepaths, no_jp2, no_xml]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The cell below defines a function for uploading files to an S3 bucket:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def upload_files(filepaths):\n",
    "    \n",
    "    # cd into the 'save' directory \n",
    "    os.chdir(OUTPUT_SAVE_DIR)\n",
    "    \n",
    "    # connects to boto3\n",
    "    s3 = boto3.resource('s3')\n",
    "    \n",
    "    # uploads files to mirror the ChronAm structure\n",
    "    for filepath in filepaths:\n",
    "        \n",
    "        if '_stats.json' in filepath or '.txt' in filepath: \n",
    "            s3.Bucket(\"S3-BUCKET-HERE\").upload_file(filepath, \"chronam_processed/\"+filepath)\n",
    "            continue\n",
    "            \n",
    "        if '.json' in filepath:\n",
    "            dir_name = filepath.split(\"/\")[0] + \"/\"\n",
    "            new_filepath = dir_name + \"data/\" + filepath.split(\"data_\")[1].replace(dir_name, '').replace('_', '/')\n",
    "            s3.Bucket(\"S3-BUCKET-HERE\").upload_file(filepath, \"chronam_processed/\"+new_filepath)\n",
    "            continue\n",
    "\n",
    "        if '.jpg' in filepath:\n",
    "            dir_name = filepath.split(\"/\")[0] + \"/\"\n",
    "            new_filepath = dir_name + \"data/\" + filepath.split(\"data_\")[1].replace(dir_name, '').replace('_', '/') \n",
    "            new_filepath = new_filepath[:new_filepath.rfind(\"/\")] + \"_\" + new_filepath[new_filepath.rfind(\"/\")+1:]\n",
    "            new_filepath = new_filepath[:new_filepath.rfind(\"/\")] + \"_\" + new_filepath[new_filepath.rfind(\"/\")+1:]\n",
    "            s3.Bucket(\"S3-BUCKET-HERE\").upload_file(filepath, \"chronam_processed/\"+new_filepath)\n",
    "            continue\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The cell below contains main(), the driver of the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from multiprocessing import Pool, get_context, Process, set_start_method\n",
    "from collections import ChainMap\n",
    "import shutil\n",
    "import time\n",
    "\n",
    "# need main for setting multiprocessing start method to spawn\n",
    "if __name__ == '__main__':\n",
    "    \n",
    "    # sets directory location where the notebook is\n",
    "    NOTEBOOK_DIR = os.getcwd()\n",
    "    os.chdir('../')\n",
    "    # sets destination for saving downloaded S3 files\n",
    "    S3_SAVE_DIR = os.getcwd() + '/chronam_files/'\n",
    "    # sets destination for output files, containing new metadata\n",
    "    OUTPUT_SAVE_DIR = os.getcwd() + '/chronam_output/'\n",
    "    os.chdir('notebooks/')\n",
    "\n",
    "    # construct the directories\n",
    "    if not os.path.isdir(S3_SAVE_DIR):\n",
    "        os.mkdir(S3_SAVE_DIR)\n",
    "    if not os.path.isdir(OUTPUT_SAVE_DIR):\n",
    "        os.mkdir(OUTPUT_SAVE_DIR)\n",
    "\n",
    "    # sets batch size for GPU inference; using 1 to prevent any memory overflow edge case in pipeline\n",
    "    INFERENCE_BATCH_SIZE = 1\n",
    "\n",
    "    # sets number of processes (be careful based on number of available cores)\n",
    "    N_CPU_PROCESSES = 48\n",
    "\n",
    "    # sets number of GPUs available\n",
    "    N_GPUS = torch.cuda.device_count()\n",
    "\n",
    "    # sets multiprocessing pool\n",
    "    pool = Pool(N_CPU_PROCESSES)  \n",
    "\n",
    "    # sets start method to spawn for GPU multiprocessing\n",
    "    ctx = get_context('forkserver')\n",
    "\n",
    "    # grabs all of the manifests\n",
    "    manifests = glob.glob(\"../manifests/*.txt\")\n",
    "\n",
    "    # now we iterate over all of the manifests\n",
    "    for manifest in manifests:\n",
    "            \n",
    "        # sets directory name\n",
    "        dir_name = manifest.split('/')[-1][:-4] + \"/\"\n",
    "\n",
    "        # first, we make the subdirectories for this manifest\n",
    "        if not os.path.isdir(S3_SAVE_DIR + dir_name):\n",
    "            os.mkdir(S3_SAVE_DIR + dir_name)\n",
    "        if not os.path.isdir(OUTPUT_SAVE_DIR + dir_name):\n",
    "            os.mkdir(OUTPUT_SAVE_DIR + dir_name)\n",
    "\n",
    "        # read manifest\n",
    "        page_filepaths = open(manifest, \"r\").read().split('\\n')\n",
    "        \n",
    "        # remove duplicate filepaths\n",
    "        page_filepaths = collections.Counter(page_filepaths).keys()\n",
    "        \n",
    "        # remove empty strings\n",
    "        filtered_page_filepaths = []\n",
    "        for filepath in page_filepaths:\n",
    "            if filepath != '':\n",
    "                filtered_page_filepaths.append(filepath)\n",
    "        page_filepaths = filtered_page_filepaths\n",
    "        \n",
    "        # if there are no files in the manifest, we skip over this newspaper manifest\n",
    "        if len(page_filepaths) == 0 or page_filepaths == ['']:\n",
    "            print(manifest)\n",
    "\n",
    "        print(\"PROCESSING MANIFEST: \" + str(dir_name) + \" (\" + str(len(page_filepaths)) + \" files)\")\n",
    "\n",
    "        print(\"validating filepaths...\")\n",
    "\n",
    "        # we check to ensure that all of these files exist; if some don't, we save the filepaths separately from the\n",
    "        # main execution path\n",
    "        packed_list = pool.map(files_exist, chunk(page_filepaths, N_CPU_PROCESSES))\n",
    "\n",
    "        good_filepaths = []\n",
    "        no_jp2 = []\n",
    "        no_xml = []\n",
    "        # we now unroll the lists from the different processes\n",
    "        for contents in packed_list:\n",
    "            good_filepaths.extend(contents[0])\n",
    "            no_jp2.extend(contents[1])\n",
    "            no_xml.extend(contents[2])\n",
    "\n",
    "        # make sure all of the files have been tested\n",
    "        assert len(page_filepaths) == len(good_filepaths) + len(no_jp2) + len(no_xml)\n",
    "\n",
    "        # we now write this info to files for quick summarization\n",
    "        with open(OUTPUT_SAVE_DIR + dir_name + 'processed_filepaths.txt', 'w') as f:\n",
    "            for filepath in good_filepaths:\n",
    "                f.write(\"%s\\n\" % filepath)\n",
    "\n",
    "        with open(OUTPUT_SAVE_DIR + dir_name + 'no_jp2.txt', 'w') as f:\n",
    "            for filepath in no_jp2:\n",
    "                f.write(\"%s\\n\" % filepath)\n",
    "\n",
    "        with open(OUTPUT_SAVE_DIR + dir_name + 'no_xml.txt', 'w') as f:\n",
    "            for filepath in no_xml:\n",
    "                f.write(\"%s\\n\" % filepath)\n",
    "\n",
    "        # now we cd into the directory for the computations\n",
    "        os.chdir(S3_SAVE_DIR + dir_name)\n",
    "        \n",
    "        # runs multiprocess for downloading of files in manifest\n",
    "        print(\"retrieving files for manifest...\")\n",
    "        # chunks good filepaths for multiprocessing\n",
    "        good_filepath_chunks = chunk(good_filepaths, N_CPU_PROCESSES)\n",
    "        # adds directory to cd into (each process starts in local path of notebook)\n",
    "        for i in range(0, len(good_filepath_chunks)):\n",
    "            good_filepath_chunks[i] = [S3_SAVE_DIR + dir_name, good_filepath_chunks[i]]\n",
    "        # calls the multiprocessing\n",
    "        image_size_dicts = pool.map(retrieve_files, good_filepath_chunks)\n",
    "\n",
    "        # we now combine the dictionaries into one\n",
    "        image_size_dict = dict(ChainMap(*image_size_dicts))\n",
    "        \n",
    "        # now we generate predictions on all of the downloaded files\n",
    "        print(\"predicting on pages...\")\n",
    "        \n",
    "        # FOR MULTIPROCESSING\n",
    "        chunked_image_filepaths = chunk(glob.glob(\"*.jpg\"), N_GPUS)\n",
    "\n",
    "        # https://stackoverflow.com/questions/31386613/python-multiprocessing-what-does-process-join-do\n",
    "        processes = []\n",
    "        for i in range(0, N_GPUS):\n",
    "            zipped = [S3_SAVE_DIR, OUTPUT_SAVE_DIR, dir_name, INFERENCE_BATCH_SIZE, chunked_image_filepaths[i], i]\n",
    "            p = ctx.Process(target=generate_predictions, args=(zipped,))\n",
    "            p.start()\n",
    "            processes.append(p)\n",
    "            \n",
    "        for process in processes:\n",
    "            process.join()\n",
    "            \n",
    "        # now, we cd into the directory containing the output files\n",
    "        os.chdir(OUTPUT_SAVE_DIR + dir_name)\n",
    "\n",
    "        # now, we grab the JSON predictions and append on image width and height so the data can be zipped \n",
    "        # for multiprocessing\n",
    "        # we want to pass these to the OCR retrieval function because they are necessary to compute bounding\n",
    "        # boxes relative to METS/ALTO OCR, and opening the image using PIL or the equivalent is costly due to\n",
    "        # the latency in loading the image into memory\n",
    "        json_filepaths = glob.glob(\"*.json\")\n",
    "\n",
    "        # grabs the \n",
    "        json_info = []\n",
    "\n",
    "        for json_filepath in json_filepaths:\n",
    "            im_width, im_height = image_size_dict[json_filepath.replace('.json', '.jpg')]\n",
    "            json_info.append([json_filepath, im_width, im_height]) \n",
    "\n",
    "        chunked_json_info = chunk(json_info, N_CPU_PROCESSES)\n",
    "        for i in range(0, len(chunked_json_info)):\n",
    "            chunked_json_info[i] = [OUTPUT_SAVE_DIR, dir_name, chunked_json_info[i]]\n",
    "\n",
    "        print(\"grabbing OCR...\")\n",
    "        pool.map(retrieve_ocr, chunked_json_info) \n",
    "        \n",
    "        print(\"cropping images...\")\n",
    "        zipped = chunk(json_filepaths, N_CPU_PROCESSES)\n",
    "        for i in range(0, len(zipped)):\n",
    "            zipped[i] = [OUTPUT_SAVE_DIR, S3_SAVE_DIR, dir_name, zipped[i]]\n",
    "        pool.map(crop, zipped)\n",
    "\n",
    "        print(\"generating embeddings...\")\n",
    "        \n",
    "        # FOR MULTIPROCESSING\n",
    "        chunked_json_filepaths = chunk(json_filepaths, N_GPUS)\n",
    "\n",
    "        # https://stackoverflow.com/questions/31386613/python-multiprocessing-what-does-process-join-do\n",
    "        processes = []\n",
    "        for i in range(0, N_GPUS):\n",
    "            zipped = [OUTPUT_SAVE_DIR, S3_SAVE_DIR, dir_name, chunked_json_filepaths[i], i]\n",
    "            p = ctx.Process(target=generate_embeddings, args=(zipped,))\n",
    "            p.start()\n",
    "            processes.append(p)\n",
    "            \n",
    "        for process in processes:\n",
    "            process.join()\n",
    "\n",
    "        print(\"uploading...\")\n",
    "            \n",
    "        # now, we cd back into the 'save' directory \n",
    "        os.chdir(OUTPUT_SAVE_DIR)\n",
    "\n",
    "        # we now compute stats on the processed newspaper and save as json\n",
    "        stats = {}\n",
    "        paths = glob.glob(dir_name + \"*.json\")\n",
    "        stats[\"processed_page_ct\"] = len(paths)\n",
    "        for path in paths:\n",
    "            if \"embeddings\" in path:\n",
    "                continue\n",
    "            # loads the JSON\n",
    "            with open(path) as f:\n",
    "                data = json.load(f)\n",
    "                stats[data[\"filepath\"]] = data\n",
    "           \n",
    "        # save stats to file\n",
    "        with open(dir_name + dir_name[:-1] + \"_stats.json\", \"w\") as fp:\n",
    "            json.dump(stats, fp)\n",
    "            \n",
    "        # we write the JSON stats file to S3 bucket\n",
    "        s3 = boto3.resource('s3')\n",
    "        s3.Bucket('S3-BUCKET-HERE').upload_file(dir_name + dir_name[:-1] + \"_stats.json\", \"chronam_stats/\" + dir_name[:-1] + \"_stats.json\")\n",
    "        os.remove(dir_name + dir_name[:-1] + \"_stats.json\")\n",
    "        \n",
    "        # now, we grab all of the files and upload to the S3 bucket in parallel\n",
    "        all_paths = glob.glob(\"**/*\", recursive=True)\n",
    "        # we filter out directory paths and keep only filepaths\n",
    "        filepaths = []\n",
    "        for path in all_paths:\n",
    "            if os.path.isfile(path):\n",
    "                filepaths.append(path)\n",
    "        # we now upload in parallel\n",
    "        filepath_chunks = chunk(filepaths, N_CPU_PROCESSES)\n",
    "        pool.map(upload_files, filepath_chunks)\n",
    "        \n",
    "        os.chdir(dir_name)\n",
    "        \n",
    "        # we now remove the folder & its contents to free up disk space\n",
    "        for path in glob.glob(\"**/*.json\", recursive=True):\n",
    "            os.remove(path)\n",
    "        for path in glob.glob(\"**/*.txt\", recursive=True):\n",
    "            os.remove(path)\n",
    "        for path in glob.glob(\"**/*.jpg\", recursive=True):\n",
    "            os.remove(path)\n",
    "        os.chdir('../')\n",
    "        shutil.rmtree(os.getcwd() + \"/\" + dir_name)\n",
    "\n",
    "        # navigate to the ChronAm pages, remove them (as well as the empty folder) \n",
    "        # and navigate back to the notebook directory\n",
    "        os.chdir(S3_SAVE_DIR + dir_name)\n",
    "\n",
    "        for path in glob.glob(\"*.xml\"):\n",
    "            os.remove(path)\n",
    "        for path in glob.glob(\"*.jpg\"):\n",
    "            os.remove(path)\n",
    "\n",
    "        os.chdir('../')\n",
    "        os.rmdir(os.getcwd() + \"/\" + dir_name)\n",
    "        os.chdir(\"../notebooks/\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}