anime_screenshot_pipeline.ipynb

{"cells":[{"attachments":{},"cell_type":"markdown","metadata":{"id":"iONsIjYGNjs3"},"source":["# Anime Screenshot Pipeline\n","https://github.com/cyber-meow/anime_screenshot_pipeline\n","\n","Colab version of: Semi-automatic pipeline to extract image training set from anime for generative model training.\n","\n","<details>\n","  <summary><big>Feature</big></summary>\n","<ul>\n","  <li>Can be found at cyber-meow github repo <a href='https://github.com/cyber-meow/anime_screenshot_pipeline#table-of-contents' target='_blank'>README.md</a></li>\n","</ul>\n","</details>\n","\n","<details>\n","  <summary><big>Limitation</big></summary>\n","<ul>\n","  <li>Not all code tested</li>\n","  <li>Not all steps from github repo is implemented</li>\n","  <li>Since this notebook is combination of many steps, dependency or package conflict may occurs</li>\n","  <li>Bad english language spelling and grammar (english is not my primary language :#). Feel free to correct and Pull Request!</li>\n","</ul>\n","</details>\n","\n","<details>\n","  <summary><big>Credits</big></summary>\n","<ul>\n","  <li>Author</li>\n","  <ul>\n","    <li><a href='https://github.com/TheSkinnyRat' target='_blank'>TheSkinnyRat</a></li>\n","  </ul>\n","\n","  <li>Base Code Repo</li>\n","  <ul>\n","    <li><a href='https://github.com/cyber-meow/anime_screenshot_pipeline' target='_blank'>cyber-meow github repo</a></li>\n","  </ul>\n","\n","  <li>Colab Template and Reference</li>\n","  <ul>\n","    <li><a href='https://github.com/Linaqruf/kohya-trainer' target='_blank'>Linaqruf/kohya-trainer</a></li>\n","  </ul>\n","\n","  <li>Code Assistant</li>\n","  <ul>\n","    <li><a href='https://chat.openai.com/' target='_blank'>OpenAI ChatGPT</a></li>\n","  </ul>\n","\n","  <li>Original cyber-meow repo credits</li>\n","  <ul>\n","    <li>This is a collection of many resources found on internet (credit to the orignal authors), and some python code written by myself and ChatGPT.</li>\n","  </ul>\n","</ul>\n","</details>\n","\n","<details>\n","  <summary><big>Whats new?</big></summary>\n","<ul>\n","  <li>(02/27/23):</li>\n","  <ul>\n","    <li>Add <code>WANDB_MODE=disabled</code> to <code>[7.2] Begin Training</code></li>\n","    <li>Fix <code>cc_training_test_set</code> error in <code>[7.2] Begin Training</code></li>\n","    <li>Fix <code>fa_move_file</code> error in <code>[8.1] Only keep images with faces and resize</code></li>\n","    <li>Move <code>[6.3] Data Cleansing</code> to <code>[5.2] Data Cleansing</code> and change to <code>n_faces</code> instead <code>n_poeple</code></li>\n","  </ul>\n","\n","  <li>(02/26/23):</li>\n","  <ul>\n","    <li>Initial First Release 🎉</li>\n","  </ul>\n","</ul>\n","</details>"]},{"cell_type":"markdown","metadata":{"id":"6GvoyYarRcAu"},"source":["| Notebook Name | Link | Repo |\n","| --- | --- | --- |\n","| [Anime Screenshot Pipeline](https://github.com/TheSkinnyRat/anime_screenshot_pipeline_colab/blob/main/anime_screenshot_pipeline.ipynb) | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=for-the-badge)](https://colab.research.google.com/github/TheSkinnyRat/anime_screenshot_pipeline_colab/blob/main/anime_screenshot_pipeline.ipynb) | [![](https://img.shields.io/static/v1?message=Github&logo=github&labelColor=5c5c5c&color=0f80c1&label=%20&style=for-the-badge)](https://github.com/TheSkinnyRat/anime_screenshot_pipeline_colab) |\n"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"TPwVMMvhOA89"},"outputs":[],"source":["!nvidia-smi"]},{"cell_type":"markdown","metadata":{"id":"PTA5atrtPCBI"},"source":["# [0] Drive Settings (Optional)"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"q6m6LjjyN0h5"},"outputs":[],"source":["#@title ## [0.1] Mount Drive\n","from google.colab import drive\n","drive.mount('/content/drive')"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"AUiMvo1WN6UT"},"outputs":[],"source":["#@title ## [0.2] Open Special `File Explorer` for Colab\n","#@markdown This will work in real-time even when you run other cells\n","!pip -q install --upgrade gdown imjoy-elfinder\n","\n","import threading\n","from google.colab import output\n","from imjoy_elfinder.app import main\n","%store -r\n","\n","thread = threading.Thread(target=main, args=[[\"--root-dir=/content\", \"--port=8765\"]])\n","thread.start()\n","\n","open_in_new_tab = True #@param {type:\"boolean\"}\n","\n","if open_in_new_tab:\n","  output.serve_kernel_port_as_window(8765)\n","else:\n","  output.serve_kernel_port_as_iframe(8765, height='500')\n"]},{"cell_type":"markdown","metadata":{"id":"8n6YvLo3PMlm"},"source":["# [1] Repository Settings"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"9xQ8Nxy3PTGC"},"outputs":[],"source":["#@title ## [1.1] Clone Reposiory\n","#@markdown Clone Anime Screenshot Pipeline from GitHub.\n","import os\n","import zipfile\n","import shutil\n","\n","root_dir = \"/content\" #@param {type: \"string\"}\n","repo_url = \"https://github.com/cyber-meow/anime_screenshot_pipeline\" #@param {type: \"string\"}\n","repo_dir = os.path.join(root_dir,\"anime_screenshot_pipeline\")\n","\n","def clone_repo(url):\n","  os.chdir(root_dir)\n","  !git clone {repo_url}\n","\n","clone_repo(repo_url)"]},{"cell_type":"markdown","metadata":{"id":"oaEJb1ClSfMi"},"source":["# [2] Data Acquisition\n","\n","You have 2 options for acquiring your dataset (.mp4 video): \n","1. Uploading it to Colab's local files.\n","2. Locating your dataset from `Google Drive` or `HuggingFace`."]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"NLHUGvmjTBEF"},"outputs":[],"source":["#@title ## [2.1] Locating Data Directory\n","#@markdown Define location of your data (.mp4 video). This cell will also create a folder based on your input.\n","import os\n","%store -r\n","\n","mp4_data_dir = \"/content/mp4_data\" #@param {'type' : 'string'}\n","%store mp4_data_dir\n","\n","os.makedirs(mp4_data_dir, exist_ok = True)\n","print(f\"Your mp4 data directory : {mp4_data_dir}\")"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"OMAVPFGnUECg"},"outputs":[],"source":["#@title ## [2.2] Unzip Mp4 Video Data\n","print('installing apt dependencies')\n","!apt-get install -y unzip aria2 > /dev/null\n","import os\n","import shutil\n","from pathlib import Path\n","from IPython.utils import capture\n","%store -r\n","\n","#@markdown Specify this section if your mp4 data is in a `zip` file and has been uploaded somewhere. This will download your dataset and automatically extract it to the `mp4_data_dir` if the `unzip_to` is empty. \n","#@markdown > Get **your** huggingface `WRITE/READ` token [here](https://huggingface.co/settings/tokens)\n","zipfile_url = \"https://huggingface.co/datasets/TheSkinnyRat/majo_no_tabitabi/resolve/main/majo_no_tabitabi-720p.zip\" #@param {'type': 'string'}\n","zipfile_name = \"zipfile.zip\"\n","unzip_to = \"\" #@param {'type': 'string'}\n","\n","hf_token = 'hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXX' #@param {'type': 'string'}\n","user_header = f\"\\\"Authorization: Bearer {hf_token}\\\"\"\n","\n","if unzip_to:\n","  os.makedirs(unzip_to, exist_ok=True)\n","else:\n","  unzip_to = mp4_data_dir\n","\n","def download_dataset(url):\n","  if url.startswith(\"/content\"):\n","    !unzip -j -o {url} -d \"{mp4_data_dir}\"\n","  elif url.startswith(\"https://drive.google.com\"):\n","    os.chdir(root_dir)\n","    !gdown --fuzzy {url}\n","  elif url.startswith(\"https://huggingface.co/\"):\n","    if '/blob/' in url:\n","      url = url.replace('/blob/', '/resolve/')\n","    !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n","  else:\n","    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n","\n","download_dataset(zipfile_url)\n","\n","os.chdir(root_dir)\n","\n","if not zipfile_url.startswith(\"/content\"):\n","  !unzip -j -o \"{root_dir}/{zipfile_name}\" -d \"{unzip_to}\"\n","  os.remove(f\"{root_dir}/{zipfile_name}\")"]},{"cell_type":"markdown","metadata":{"id":"ab79RaqkgbRq"},"source":["# [3] [#](https://github.com/cyber-meow/anime_screenshot_pipeline#frame-extraction) Frame Extraction\n","\n","Extract 5000~10000 frames per episode of 24 minutes"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"YJehB7DmggGP"},"outputs":[],"source":["#@title ## [3.1] Begin Frame Extraction\n","#@markdown Define your Frame Extraction directory output\n","fe_dest_dir = \"/content/fe_output\" #@param {'type': 'string'}\n","#@markdown **ATTENTION:** Rename your mp4 video file data to the following order to to avoid unwanted errors\\\n","#@markdown Your mp4 file pattern. Example, this pattern will extract frame for following files\n","#@markdown - majo_no_tabitabi_01.mp4\n","#@markdown - majo_no_tabitabi_02.mp4\n","#@markdown - majo_no_tabitabi_03.mp4\n","#@markdown - ...\n","fe_pattern = \"majo_no_tabitabi_*.mp4\" #@param {'type': 'string'}\n","\n","os.chdir(repo_dir)\n","!python extract_frames.py --src_dir \"{mp4_data_dir}\" \\\n","--dst_dir \"{fe_dest_dir}\" \\\n","--prefix series_episode \\\n","--pattern \"{fe_pattern}\""]},{"cell_type":"markdown","metadata":{"id":"mHlskSyZF5Qe"},"source":["# [4] [#](https://github.com/cyber-meow/anime_screenshot_pipeline#similar-image-removal) Similar Image Removal\n","\n","Reduce dataset size by a factor of 10 by removing similar images"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"-PEh7GRmGCQH"},"outputs":[],"source":["#@title ## [4.1] Begin Removal\n","!pip -q install --upgrade fiftyone\n","import numpy as np\n","import fiftyone as fo\n","import fiftyone.zoo as foz\n","\n","dataset_dir = fe_dest_dir\n","dataset = fo.Dataset.from_dir(dataset_dir, dataset_type=fo.types.ImageDirectory)\n","\n","model = foz.load_zoo_model(\"mobilenet-v2-imagenet-torch\")\n","embeddings = dataset.compute_embeddings(model)\n","\n","print(embeddings.shape)\n","\n","from tqdm import tqdm\n","\n","def mark_duplicate(subdataset, similarity_matrix, thresh=0.985):\n","    \n","    n = len(similarity_matrix)\n","    similarity_matrix = similarity_matrix - np.identity(n)\n","    \n","    id_map = [s.id for s in subdataset.select_fields([\"id\"])]\n","    samples_to_remove = set()\n","    samples_to_keep = set()\n","\n","    for idx, sample in enumerate(subdataset):\n","        max_similarity = similarity_matrix[idx].max()\n","        sample[\"max_similarity\"] = max_similarity\n","        sample.save()\n","\n","    for idx, sample in tqdm(enumerate(subdataset)):\n","        if sample.id not in samples_to_remove:\n","            # Keep the first instance of two duplicates\n","            samples_to_keep.add(sample.id)\n","\n","            dup_idxs = np.where(similarity_matrix[idx] > thresh)[0]\n","            for dup in dup_idxs:\n","                # We kept the first instance so remove all other duplicates\n","                samples_to_remove.add(id_map[dup])\n","\n","            if len(dup_idxs) > 0:\n","                sample.tags.append(\"has_duplicates\")\n","                sample.save()\n","\n","        else:\n","            sample.tags.append(\"duplicate\")\n","            sample.save()\n","    return samples_to_remove, samples_to_keep\n","\n","from sklearn.metrics.pairwise import cosine_similarity\n","\n","max_compare_size = 10000\n","thresh = 0.985\n","\n","samples_to_remove = set()\n","samples_to_keep = set()\n","\n","for k in range(0, len(embeddings), max_compare_size):\n","    end = min(k + max_compare_size, len(embeddings))\n","    similarity_matrix = cosine_similarity(embeddings[k:end])\n","    samples_to_remove_sub, samples_to_keep_sub = mark_duplicate(\n","        dataset[k:end], similarity_matrix, thresh)\n","    samples_to_remove = samples_to_remove | samples_to_remove_sub\n","    samples_to_keep = samples_to_keep | samples_to_keep_sub\n","\n","sir_visualize_dataset = False #@param {'type' : 'boolean'}\n","if sir_visualize_dataset:\n","  session = fo.launch_app(dataset)\n","\n","import os\n","for sample_id in tqdm(samples_to_remove):\n","    os.remove(dataset[sample_id].filepath)\n","dataset.delete_samples(list(samples_to_remove))"]},{"cell_type":"markdown","metadata":{"id":"nbz62levLo97"},"source":["# [5] [#](https://github.com/cyber-meow/anime_screenshot_pipeline#face-detection-and-cropping) Face Detection and Cropping"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"z0AXSVnHL6C0"},"outputs":[],"source":["#@title ## [5.1] Add face information to metadata\n","\n","# ERROR ISSUE: https://github.com/hysts/anime-face-detector/issues/13#issuecomment-1419694747\n","!pip -q install numpy scipy numba --upgrade\n","\n","# Downgrading Torch\n","!pip -q install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117\n","\n","!pip -q install openmim==0.2.1\n","!pip -q install mmcv-full==1.7.0 -f https://download.openmmlab.com/mmcv/dist/cu117/torch1.13/index.html\n","!pip -q install mmdet==2.25.1\n","!pip -q install mmpose==0.28.1\n","!pip -q install moviepy==1.0.3\n","#!git clone https://github.com/hysts/anime-face-detector\n","!pip -q install anime-face-detector\n","\n","os.chdir(repo_dir)\n","!python detect_faces.py --src_dir \"{fe_dest_dir}\""]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["#@title ## [5.2] Data Cleansing\n","#@markdown This will delete the dataset and their metadata if `n_faces==0` in the `.json` file metadata\n","import os\n","import json\n","\n","# Define the directory path\n","at_cleansing_path = fe_dest_dir\n","\n","# Traverse the directory tree\n","for root, dirs, files in os.walk(at_cleansing_path):\n","    for file in files:\n","        # Check if the file is a .json file\n","        if file.endswith('.json'):\n","            # Construct the full file path\n","            file_path = os.path.join(root, file)\n","            # Open the JSON file\n","            with open(file_path, 'r') as f:\n","                json_data = json.load(f)\n","            # Check if the 'n_people' field is 0\n","            if json_data['n_faces'] == 0:\n","                # Construct the image file path and delete the file\n","                image_file_path = os.path.splitext(file_path)[0] + '.png'\n","                if os.path.exists(image_file_path):\n","                    os.remove(image_file_path)\n","                # Delete the .json and .tags files\n","                os.remove(file_path)\n","                tags_file_path = os.path.splitext(file_path)[0] + '.png.tags'\n","                if os.path.exists(tags_file_path):\n","                    os.remove(tags_file_path)\n","\n","print('Done! Data Cleaned')"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"CVTueXJyM-uR"},"outputs":[],"source":["#@title ## [5.3] Crop out the maximum square for each face\n","#@markdown Indicates the minimum number of faces that should be contained in the original image for the cropping to take place. Note that the cropped images are store in the same folders as the input images and no cropping is performed for face that is too large\n","fdc_min_face_number = 2 #@param {type :\"integer\"}\n","\n","os.chdir(repo_dir)\n","!python crop_faces.py --src_dir \"{fe_dest_dir}\" --min_face_number {fdc_min_face_number}"]},{"cell_type":"markdown","metadata":{"id":"PhX3sLNiXifa"},"source":["# [6] [#](https://github.com/cyber-meow/anime_screenshot_pipeline#automatic-tagging) Automatic Tagging"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"xiQu6HODXri5"},"outputs":[],"source":["#@title ## [6.1] Tag your images with an off-the-shelf tagger\n","!pip -q install \"tensorflow<2.11\"\n","!pip -q install huggingface-hub\n","\n","os.chdir(os.path.join(repo_dir, \"tagger\"))\n","!python tag_images_by_wd14_tagger.py --batch_size 16 --caption_extension \".tags\" \"{fe_dest_dir}\""]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"lulAxkgNYlPf"},"outputs":[],"source":["#@title ## [6.2] Save tag information into metadata\n","os.chdir(repo_dir)\n","!python augment_metadata.py --use_tags --tag_extension \"tags\" --general_description \"aniscreen\" --src_dir \"{fe_dest_dir}\""]},{"cell_type":"markdown","metadata":{"id":"Q0Leaj0XZ5bH"},"source":["# [7] [#](https://github.com/cyber-meow/anime_screenshot_pipeline#character-classification-with-few-shot-learning) Character Classification with Few-Shot Learning\n","Train your own model for series-specific concepts"]},{"cell_type":"markdown","metadata":{"id":"OxvE96gUnB6y"},"source":["## [7.1] Dataset Preparation\n","This is a simple classification task. To begin just create a directory `training_data_original` and organize the directory as following\n","```\n","├── ./elaina_(majo_no_tabitabi)\n","├── ./saya_(majo_no_tabitabi)\n","├── ./character3\n","├── ./character4\n","...\n","└── ./ood\n","```\n","Put a few images in each folder that try to capture different variations of the character. In my test 10~20 is good enough. Zip the whole `training_data_original` folder and upload to hugging face or somewhere else, or you can directly upload to colab notebook using file manager.\n","> Note: [Original Instruction](https://github.com/cyber-meow/anime_screenshot_pipeline#dataset-preparation)\n","\n","> If you already have `.cktp` training file, you can skip `[7.1]` and `[7.2]` step and go to `[7.3-C]` to define and locate your `.ckpt` file"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":39407,"status":"ok","timestamp":1677409632513,"user":{"displayName":"TheSkinnyRat AI","userId":"03355026868987494301"},"user_tz":-420},"id":"s8QnUPe3b3-N","outputId":"f699e875-6c62-473f-8160-a9df5c626cd7"},"outputs":[],"source":["os.chdir(os.path.join(repo_dir, \"classifier_training\"))\n","!pip -q install -r requirements.txt\n","os.chdir(os.path.join(repo_dir, \"classifier_training/models\"))\n","!pip -q install -e .\n","\n","cc_zipfile_url = \"https://huggingface.co/datasets/TheSkinnyRat/majo_no_tabitabi/resolve/main/training_data_original.zip\" #@param {'type': 'string'}\n","cc_zipfile_name = \"zipfile.zip\"\n","cc_unzip_to = \"\" #@param {'type': 'string'}\n","#@markdown If you want to split into training and test set you can check this.\\\n","#@markdown It does a 70%/30% train/test split. \n","#@markdown > Note: Not tested yet\n","cc_split_training_and_test = False #@param {'type': 'boolean'}\n","\n","user_header = f\"\\\"Authorization: Bearer {hf_token}\\\"\"\n","\n","if cc_unzip_to:\n","  os.makedirs(cc_unzip_to, exist_ok=True)\n","else:\n","  cc_unzip_to = \"/content/\"\n","\n","def download_dataset(url):\n","  if url.startswith(\"/content\"):\n","    !unzip -j -o {url} -d \"/content/\"\n","  elif url.startswith(\"https://drive.google.com\"):\n","    os.chdir(root_dir)\n","    !gdown --fuzzy {url}\n","  elif url.startswith(\"https://huggingface.co/\"):\n","    if '/blob/' in url:\n","      url = url.replace('/blob/', '/resolve/')\n","    !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n","  else:\n","    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n","\n","download_dataset(cc_zipfile_url)\n","\n","os.chdir(root_dir)\n","\n","if not zipfile_url.startswith(\"/content\"):\n","  !unzip -o \"{root_dir}/{cc_zipfile_name}\" -d \"{cc_unzip_to}\"\n","  os.remove(f\"{root_dir}/{cc_zipfile_name}\")\n","\n","cc_src_dir = os.path.join(cc_unzip_to, \"training_data_original\")\n","cc_dataset_dir = os.path.join(cc_unzip_to, \"training_dataset\")\n","\n","# replace filename contain space with underscore and lowercase\n","[os.rename(root + os.sep + file, root + os.sep + file.lower().replace(\" \", \"_\")) for root, _, files in os.walk(cc_src_dir) for file in files]\n","\n","os.chdir(os.path.join(repo_dir, \"classifier_dataset_preparation\"))\n","!python crop_and_make_dataset.py --src_dir \"{cc_src_dir}\" --dst_dir \"{cc_dataset_dir}/data\"\n","\n","os.chdir(os.path.join(repo_dir, \"classifier_dataset_preparation\"))\n","!python make_data_dic_imagenetsyle.py \"{cc_dataset_dir}\"\n","\n","if cc_split_training_and_test:\n","  os.chdir(os.path.join(repo_dir, \"classifier_dataset_preparation\"))\n","  cc_dataset_labels = os.path.join(cc_dataset_dir, \"labels.csv\")\n","  !python data_split.py \"{cc_dataset_labels}\" 0.7 0.3"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"Fkb3hi33slnX"},"outputs":[],"source":["#@title ## [7.2] Begin Training\n","!wandb disabled\n","!apt-get install -y wget > /dev/null\n","cc_checkpoint_dir = os.path.join(root_dir, \"training_model\")\n","cc_checkpoint_path = os.path.join(root_dir, \"training_model/danbooruFaces_L_16_image128_batch16_SGDlr0.001_ptTrue_seed0_warmupCosine_interTrue_mmFalse_textLenNone_maskNoneconstanttagtokenizingshufFalse_lastEpoch.ckpt\")\n","cc_trained_checkpoint_dir = os.path.join(root_dir, \"trained_model\")\n","cc_trained_checkpoint_path = os.path.join(root_dir, \"trained_model/*.ckpt\")\n","\n","if not os.path.exists(cc_checkpoint_path):\n","  if not os.path.exists(cc_checkpoint_dir):\n","    os.makedirs(cc_checkpoint_dir, exist_ok=True)\n","  os.chdir(cc_checkpoint_dir)\n","  !wget https://huggingface.co/TheSkinnyRat/public-backup/resolve/main/danbooruFaces_L_16_image128_batch16_SGDlr0.001_ptTrue_seed0_warmupCosine_interTrue_mmFalse_textLenNone_maskNoneconstanttagtokenizingshufFalse_lastEpoch.ckpt\n","#@markdown Note that testing is disabled by default. Check this if you want to use test.\\\n","#@markdown Validation will then be performed every 5 epochs (default).\n","cc_training_test_set = False #@param {'type': 'boolean'}\n","\n","#@markdown > ~~**WARNING:** Don't know what happen, but you'll get wandb requires interactive prompt like:~~ (probably fixed)\n","#@markdown ```\n","#@markdown wandb: (1) Create a W&B account\n","#@markdown wandb: (2) Use an existing W&B account\n","#@markdown wandb: (3) Don't visualize my results\n","#@markdown wandb: Enter your choice: Traceback (most recent call last): ...\n","#@markdown ```\n","#@markdown I'll only choose `(3) Don't visualize my results` by typing `3` and `enter` in the cursor result, or you want something experimental?.\n","\n","!pip install keras==2.10.0\n","# working commit https://github.com/arogozhnikov/einops/commit/f569905f8ba2f55393164262c5e4200a8cdb57ab\n","!pip install git+https://github.com/arogozhnikov/einops.git\n","\n","os.chdir(os.path.join(repo_dir, \"classifier_training\"))\n","if not cc_training_test_set:\n","  !python train.py --transfer_learning --model_name L_16 --interm_features_fc \\\n","  --batch_size=8 --no_epochs 40 --dataset_path \"{cc_dataset_dir}\" \\\n","  --results_dir \"{cc_trained_checkpoint_dir}\" \\\n","  --checkpoint_path \"{cc_checkpoint_path}\"\n","else:\n","  !python train.py --transfer_learning --model_name L_16 --interm_features_fc \\\n","  --batch_size=8 --no_epochs 40 --dataset_path \"{cc_dataset_dir}\" \\\n","  --results_dir \"{cc_trained_checkpoint_dir}\" \\\n","  --checkpoint_path \"{cc_checkpoint_path}\" \\\n","  --use_test_set"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"C8XEmx72EaI-"},"outputs":[],"source":["#@title ## [7.3-C] Define and Locate Model (`.ckpt` model required)\n","!apt-get install -y wget > /dev/null\n","cc_model_url = \"example.com/model.ckpt\" #@param {'type': 'string'}\n","cc_trained_checkpoint_dir = os.path.join(root_dir, \"trained_model\")\n","cc_trained_checkpoint_path = os.path.join(root_dir, \"trained_model/*.ckpt\")\n","\n","os.chdir(cc_trained_checkpoint_dir)\n","!wget \"{cc_model_url}\"\n","\n","print(\"Done! Model Located\")"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"EqBL1N5c0DsD"},"outputs":[],"source":["#@title ## [7.4] Inference\n","#@markdown Threshold of confidence to classify as character\n","cc_cls_thresh = 0.25 #@param {type:\"slider\", min:0.1, max:1, step:0.01}\n","#@markdown Threshold of confidence to add a tag\n","cc_tagger_thresh = 0.35 #@param {type:\"slider\", min:0.1, max:1, step:0.01}\n","#@markdown At this point we finally get a model to classify the characters of the series!\\\n","#@markdown We then write the character information into the metadata file if such information is not present.\\\n","#@markdown To overwrite in all the cases check this\n","cc_inference_overwrite = False #@param {type:\"boolean\"}\n","import glob\n","cc_trained_checkpoint_path = glob.glob('/content/trained_model/*.ckpt')[0]\n","\n","os.chdir(repo_dir)\n","if not cc_inference_overwrite:\n","  !python classify_characters.py --dataset_path \"{cc_dataset_dir}\" \\\n","  --checkpoint_path \"{cc_trained_checkpoint_path}\" \\\n","  --cls_thresh \"{cc_cls_thresh}\" --tagger_thresh \"{cc_tagger_thresh}\"\\\n","  --src_dir \"{fe_dest_dir}\"\n","else:\n","  !python classify_characters.py --dataset_path \"{cc_dataset_dir}\" \\\n","  --checkpoint_path \"{cc_trained_checkpoint_path}\" \\\n","  --cls_thresh \"{cc_cls_thresh}\" --tagger_thresh \"{cc_tagger_thresh}\"\\\n","  --src_dir \"{fe_dest_dir}\" \\\n","  --overwrite\n"]},{"cell_type":"markdown","metadata":{"id":"iHtuBbrt1xl4"},"source":["# [8] [#](https://github.com/cyber-meow/anime_screenshot_pipeline#folder-arrangement) Folder Arrangement"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"AJpUfWM_146E"},"outputs":[],"source":["#@title ## [8.1] Only keep images with faces and resize\n","#@markdown The following command can be run before cropping and tagging to eliminate images with no faces.\n","#@markdown - With `min_face_number` and `max_face_number` it only saves images whose face number is within this range to `dst_dir`.\n","#@markdown - `max_image_size` makes sure that saved images are resized so that both its width and height are smaller than the provided value.\n","#@markdown - Check `move_file` if you want to move file to destination directory instead of creating new ones. `max_image_size` is ignored in this case.\n","fa_dest_dir = \"/content/fe_output_arranged\" #@param {'type': 'string'}\n","fa_min_face = 1 #@param {'type': 'integer'}\n","fa_max_face = 10 #@param {'type': 'integer'}\n","fa_max_image_size = 1024 #@param {'type': 'integer'}\n","fa_move_file = False #@param {'type': 'boolean'}\n","\n","os.chdir(repo_dir)\n","if not fa_move_file:\n","  !python arrange_folder.py --min_face_number {fa_min_face} --max_face_number {fa_max_face} \\\n","  --keep_src_structure --format '' --max_image_size {fa_max_image_size} \\\n","  --src_dir \"{fe_dest_dir}\" --dst_dir \"{fa_dest_dir}\"\n","else:\n","  !python arrange_folder.py --min_face_number {fa_min_face} --max_face_number {fa_max_face} \\\n","  --keep_src_structure --format '' --max_image_size {fa_max_image_size} \\\n","  --src_dir \"{fe_dest_dir}\" --dst_dir \"{fa_dest_dir}\" \\\n","  --move_file"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"mOZvW7Ev5UvT"},"outputs":[],"source":["#@title ## [8.2] Arrange the folder in hierarchy using metadata\n","#@markdown The folder structure itself is specified by the argument `format`. Different levels of folders are separated by `/`. Accepted folder types are `n_characters`, `n_faces`, `n_people`, `character`, and `fh_ratio`.\n","#@markdown - `n_characters`: This creates folders using the number of characters. Passing argument `max_character_number` = `6` puts all the scenes with more than 6 characters into the folder `6+_characters`.\n","#@markdown - `character`: This creates folders with sorted character names split by `+`. To avoid creating a specific folder for character combination that shows up too few times, we pass the argument `min_image_per_combination` so that images of all the character combinations with fewer than a certain number of images are saved in `.../character_others`.\n","#@markdown - `fh_ratio`: This creates folders according to the maximum face height ratio.\n","\n","fa_heirarchy_dest_dir = \"/content/fe_output_arranged_hierarchy\" #@param {'type': 'string'}\n","fa_format = \"n_characters/character/fh_ratio\" #@param {'type': 'string'}\n","fa_max_character_number = 6 #@param {'type': 'integer'}\n","fa_min_image_per_combination = 10 #@param {'type': 'integer'}\n","\n","os.chdir(repo_dir)\n","!python arrange_folder.py \\\n","--move_file --format \"{fa_format}\" \\\n","--max_character_number {fa_max_character_number} --min_image_per_combination {fa_min_image_per_combination} \\\n","--src_dir \"{fa_dest_dir}\" --dst_dir \"{fa_heirarchy_dest_dir}\""]},{"cell_type":"markdown","metadata":{"id":"srr3_7PHD1ZZ"},"source":["# [9] Create `.txt` File (Optional)\n"," Create `.txt` file next to every image with json metadata"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"uKfL1KwwEGbn"},"outputs":[],"source":["#@title ## [9.1] Begin Create `.txt` File\n","import os\n","import json\n","\n","#@markdown Define the directory containing the image files and their corresponding JSON files\n","ctf_dir_path = '/content/fe_output_arranged_hierarchy' #@param {type: 'string'}\n","\n","# Iterate through all files and subdirectories in the directory\n","for root, dirs, files in os.walk(ctf_dir_path):\n","    for file in files:\n","        # Check if the file is an image file (has extension .png, .jpg, or .jpeg)\n","        if file.endswith('.png') or file.endswith('.jpg') or file.endswith('.jpeg'):\n","            # Construct the path to the JSON file corresponding to this image file\n","            json_file = os.path.join(root, file.split('.')[0] + '.json')\n","            # Check if the JSON file exists\n","            if os.path.exists(json_file):\n","                # Load the JSON data from the file\n","                with open(json_file, 'r') as f:\n","                    json_data = json.load(f)\n","                # Extract the desired values from the JSON data\n","                characters = ', '.join(json_data['characters']).replace('_', ' ')\n","                general = json_data['general'].replace('_', ' ')\n","                tags = ', '.join(json_data['tags']).replace('_', ' ')\n","                # Construct the text to write to the .txt file\n","                text = f\"{characters}, {general}, {tags}\"\n","                # Construct the path to the .txt file to write to\n","                txt_file = os.path.join(root, file.split('.')[0] + '.txt')\n","                # Write the text to the .txt file\n","                with open(txt_file, 'w') as f:\n","                    f.write(text)\n","\n","print('Done! .txt File Created')\n"]},{"cell_type":"markdown","metadata":{"id":"QbuwZQdM6rim"},"source":["# [10] HuggingFace Upload (Optional)\n"," Upload your dataset to hugging face"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"VRnHNAsA06dA"},"outputs":[],"source":["#@title ### [10.1] Upload Config\n","!pip -q install --upgrade huggingface-hub\n","from huggingface_hub import login\n","from huggingface_hub import HfApi\n","from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n","\n","login(hf_token, add_to_git_credential=True)\n","\n","api = HfApi()\n","user = api.whoami(hf_token)\n","\n","#@markdown Fill this if you want to upload to your organization, or just leave it empty.\n","\n","orgs_name = \"\" #@param{type:\"string\"}\n","\n","#@markdown If your model/dataset repo didn't exist, it will automatically create your repo.\n","dataset_name = \"majo_no_tabitabi\" #@param{type:\"string\"}\n","make_this_model_private = True #@param{type:\"boolean\"}\n","\n","if orgs_name == \"\":\n","  datasets_repo = user['name']+\"/\"+dataset_name.strip()\n","else:\n","  datasets_repo = orgs_name+\"/\"+dataset_name.strip()\n","\n","if dataset_name != \"\":\n","  try:\n","      validate_repo_id(datasets_repo)\n","      api.create_repo(repo_id=datasets_repo,\n","                      repo_type=\"dataset\",\n","                      private=make_this_model_private)\n","      print(\"Dataset Repo didn't exists, creating repo\")\n","      print(\"Dataset Repo\",datasets_repo,\"created!\\n\")\n","\n","  except HfHubHTTPError as e:\n","      print(f\"Dataset repo: {datasets_repo} exists, skipping create repo\\n\")\n"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"vxmoMwl-AMHx"},"outputs":[],"source":["#@title ### [10.2] Zip Dataset\n","hu_dest_dir = \"/content/fe_output_arranged_hierarchy\" #@param {'type': 'string'}\n","hu_zip_file_name = \"final_dataset_ep1\" #@param {'type': 'string'}\n","hu_zip_path = \"/content/\"+hu_zip_file_name+\".zip\"\n","\n","os.chdir(hu_dest_dir)\n","!7z a \"{hu_zip_path}\" ./ -y\n","print(\"Done!\")"]},{"cell_type":"code","execution_count":null,"metadata":{"cellView":"form","id":"cuGdl3dR7Cov"},"outputs":[],"source":["#@title ### [10.3] Begin Upload\n","from huggingface_hub import HfApi\n","from pathlib import Path\n","import shutil\n","import zipfile\n","import os\n","\n","api = HfApi()\n","\n","#@markdown This will be compressed your `fe_dest_dir` to zip and  uploaded to datasets repo\n","\n","#@markdown  Other Information\n","commit_message = \"\" #@param {type :\"string\"}\n","temp_dataset = hu_zip_path\n","\n","if not commit_message:\n","  commit_message = f\"feat: upload {hu_zip_file_name}.zip\"\n","\n","def upload_dataset(dataset_paths):\n","  path_obj = Path(dataset_paths)\n","  #dataset_name = path_obj.parts[-1]\n","  dataset_name = f\"{hu_zip_file_name}.zip\"\n","\n","  print(f\"Uploading {dataset_name} to https://huggingface.co/datasets/\"+datasets_repo)\n","  print(f\"Please wait...\")\n","\n","  api.upload_file(\n","      path_or_fileobj=dataset_paths,\n","      path_in_repo=dataset_name,\n","      repo_id=datasets_repo,\n","      repo_type=\"dataset\",\n","      commit_message=commit_message,\n","  )\n","  print(f\"Upload success, located at https://huggingface.co/datasets/\"+datasets_repo+\"/blob/main/\"+dataset_name+\"\\n\")\n","\n","def zip_file(tmp_folders):\n","    zipfiles = temp_dataset \n","    with zipfile.ZipFile(zipfiles, 'w') as zip:\n","      for tmp_folders, dirs, files in os.walk(tmp_folders):\n","          for file in files:\n","              zip.write(os.path.join(tmp_folders, file))\n","\n","def upload():\n","  # zip_file(fe_dest_dir)\n","  upload_dataset(temp_dataset)\n","  os.remove(temp_dataset)\n","\n","upload()"]}],"metadata":{"accelerator":"GPU","colab":{"provenance":[],"toc_visible":true},"gpuClass":"standard","kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.7"},"vscode":{"interpreter":{"hash":"c096e7fa74c2db2b63cd23c0160cc87b3acf12d5f9115c8cce3b0c883f29c7bc"}}},"nbformat":4,"nbformat_minor":0}