better-captions-for-the-data-already-downloaded.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0c3b445a-9e8c-4620-9848-ca7837f704e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2023-10-10 10:33:54--  https://raw.githubusercontent.com/opencontext/archaeology-images-ai/main/json_data/artifact_images_w_sentence_captions.json\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.111.133, 185.199.108.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 62284991 (59M) [text/plain]\n",
      "Saving to: ‘artifact_images_w_sentence_captions.json’\n",
      "\n",
      "artifact_images_w_s 100%[===================>]  59.40M  6.31MB/s    in 11s     \n",
      "\n",
      "2023-10-10 10:34:06 (5.25 MB/s) - ‘artifact_images_w_sentence_captions.json’ saved [62284991/62284991]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://raw.githubusercontent.com/opencontext/archaeology-images-ai/main/json_data/artifact_images_w_sentence_captions.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cbbe5c6d-721d-4a6f-a7ac-53d0cd96caf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "## this code is just to append better captions to data that has already been downloaded\n",
    "## with the poorer captions\n",
    "\n",
    "import json\n",
    "import os\n",
    "\n",
    "def process_chunk(artifact_dict):\n",
    "    for artifact in artifact_dict:\n",
    "        for train in train_lines:\n",
    "            # Split the filename and the extension so that we only compare the filename\n",
    "            if os.path.splitext(train['image'].split('/')[-1])[0] == artifact['media__uuid']:\n",
    "                train['caption'] = artifact['caption']\n",
    "    return train_lines\n",
    "\n",
    "artifact_file = 'artifact_images_w_sentence_captions.json'\n",
    "train_file = 'train.json'\n",
    "batchsize = 1000 # Customize this as needed\n",
    "\n",
    "train_lines = [json.loads(line) for line in open(train_file)]\n",
    "\n",
    "with open(artifact_file) as f:\n",
    "    artifact_dict = json.load(f)  # Changed this line to json.load()\n",
    "\n",
    "    # Process in chunks\n",
    "    for i in range(0, len(artifact_dict), batchsize):\n",
    "        artifact_chunk = artifact_dict[i:i+batchsize]\n",
    "        train_lines = process_chunk(artifact_chunk)\n",
    "\n",
    "with open(train_file, 'w') as f:\n",
    "    for item in train_lines:\n",
    "        f.write(\"%s\\n\" % json.dumps(item))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8fd5b0a3-6061-4865-8188-b97fe09089d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import requests\n",
    "import pandas as pd\n",
    "from urllib.parse import urlparse\n",
    "from urllib.request import urlretrieve\n",
    "from sklearn.model_selection import train_test_split\n",
    "import concurrent.futures\n",
    "\n",
    "# Load JSON from remote URL\n",
    "url = \"https://raw.githubusercontent.com/opencontext/archaeology-images-ai/main/json_data/artifact_images_w_sentence_captions.json\"\n",
    "response = requests.get(url)\n",
    "data = response.json()\n",
    "\n",
    "# Randomly select records\n",
    "df = pd.DataFrame(data)\n",
    "train_df, rem_df = train_test_split(df, train_size=10, random_state=42)\n",
    "test_df = rem_df.sample(2, random_state=42)\n",
    "\n",
    "# Download and rename images\n",
    "def download_and_rename(row, folder):\n",
    "    os.makedirs(folder, exist_ok=True)  \n",
    "    uri = row['image_file__uri']\n",
    "    uuid = row['media__uuid']\n",
    "    caption = row['caption']\n",
    "    parse_object = urlparse(uri)\n",
    "    _, ext = os.path.splitext(parse_object.path)\n",
    "    new_image_name = uuid + ext\n",
    "    new_image_path = os.path.join(folder, new_image_name)\n",
    "    urlretrieve(uri, new_image_path)\n",
    "    return {\"image\": new_image_path, \"caption\": caption}\n",
    "\n",
    "# Writing to 'jsonl' files\n",
    "def write_to_jsonl(new_data, jsonl_file):\n",
    "    with open(jsonl_file, 'w') as file:\n",
    "        for json_dict in new_data:\n",
    "            line = json.dumps(json_dict)\n",
    "            file.write(line + \"\\n\")\n",
    "\n",
    "# Process train and test data\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    train_data = list(executor.map(download_and_rename, [row for _, row in train_df.iterrows()], ['ttrain_images']*len(train_df)))\n",
    "    test_data = list(executor.map(download_and_rename, [row for _, row in test_df.iterrows()], ['ttest_images']*len(test_df)))\n",
    "\n",
    "# Write train/test data to jsonl files\n",
    "write_to_jsonl(train_data, 'ttrain.json')\n",
    "write_to_jsonl(test_data, 'ttest.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87663be7-5dc6-441a-b0a8-7b6daf3fff69",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}