From d9915808db998b740655b70cd3e3a90b114a6c30 Mon Sep 17 00:00:00 2001 From: 0xlws <87901794+0xlws@users.noreply.github.com> Date: Tue, 25 Apr 2023 23:36:40 +0200 Subject: [PATCH] feat: add automated identifier renaming --- fast-DreamBooth.ipynb | 111 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 90 insertions(+), 21 deletions(-) diff --git a/fast-DreamBooth.ipynb b/fast-DreamBooth.ipynb index ddce4857..83bffd10 100644 --- a/fast-DreamBooth.ipynb +++ b/fast-DreamBooth.ipynb @@ -515,13 +515,7 @@ " if MODEL_NAME==\"\":\n", " print('\u001b[1;31mNo model found, use the \"Model Download\" cell to download a model.')\n", " else:\n", - " print('\u001b[1;32mSession created, proceed to uploading instance images')\n", - "\n", - " #@markdown\n", - "\n", - " #@markdown # The most important step is to rename the instance pictures of each subject to a unique unknown identifier, example :\n", - " #@markdown - If you have 10 pictures of yourself, simply select them all and rename only one to the chosen identifier for example : phtmejhn, the files would be : phtmejhn (1).jpg, phtmejhn (2).png ....etc then upload them, do the same for other people or objects with a different identifier, and that's it.\n", - " #@markdown - Checkout this example : https://i.imgur.com/d2lD3rz.jpeg" + " print('\u001b[1;32mSession created, proceed to uploading instance images')" ] }, { @@ -529,12 +523,27 @@ "execution_count": null, "metadata": { "cellView": "form", - "id": "LC4ukG60fgMy" + "id": "txBgZGJ0TNpZ" }, "outputs": [], "source": [ - "import shutil\n", + "#@markdown You can now use one of the two options to rename the instance pictures of each subject to a unique unknown identifier:\n", + "\n", + "#@markdown ### Option 1: Provide the root folder path containing identifier folders\n", + "#@markdown - Provide the path to the root folder containing identifier folders. The script will take care of renaming the images in each folder with their respective identifiers.\n", + "#@markdown - Example: If your root folder is \"my_instances\" and contains folders like \"phtmejhn\" and \"abzjrmcn\", the script will automatically rename images in each of those folders.\n", + "\n", + "#@markdown ### Option 2: Manually enter the identifiers\n", + "#@markdown - If you don't provide a root folder path, enter the identifiers comma-separated (e.g., phtmejhn, abzjrmcn). The script will prompt you to either provide a path to a directory containing the images for the identifier or leave the input field empty and press ENTER to choose the file upload option.\n", + "#@markdown - Example: phtmejhn (1).jpg, phtmejhn (2).png, abzjrmcn (1).jpeg, etc.\n", + "\n", + "# Import necessary libraries\n", + "import os\n", + "from glob import glob\n", "from google.colab import files\n", + "from pathlib import Path\n", + "import shutil\n", + "from shutil import copyfile\n", "import time\n", "from PIL import Image\n", "from tqdm import tqdm\n", @@ -548,14 +557,10 @@ " wget.download('https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/Dreambooth/smart_crop.py')\n", " from smart_crop import *\n", "\n", - "#@markdown #Instance Images\n", - "#@markdown ----\n", "\n", - "#@markdown\n", - "#@markdown - Run the cell to upload the instance pictures.\n", - "#@markdown - You can add `external captions` in txt files by simply giving each txt file the same name as the instance image, for example dikgur (1).jpg and dikgur (1).txt, and upload them here, to use the external captions, check the box \"external_captions\" in the training cell. `All the images must have one same extension` jpg or png or....etc\n", + "#@markdown # Instance image upload\n", "\n", - "Remove_existing_instance_images= True #@param{type: 'boolean'}\n", + "Remove_existing_instance_images= False #@param{type: 'boolean'}\n", "#@markdown - Uncheck the box to keep the existing instance images.\n", "\n", "if Remove_existing_instance_images:\n", @@ -573,9 +578,77 @@ " %rm -r $INSTANCE_DIR\"/.ipynb_checkpoints\"\n", "\n", "\n", - "IMAGES_FOLDER_OPTIONAL=\"\" #@param{type: 'string'}\n", + "#@markdown ___\n", + "\n", + "# Enter the path to the root folder containing identifier folders (empty if not using Option 1)\n", + "root_folder_path = \"\" #@param {type:\"string\"}\n", + "#@markdown - Enter the path to the root folder containing identifier folders (empty if not using this option)\n", + "\n", + "\n", + "#@markdown Or\n", + "\n", + "# Enter the identifiers (comma-separated)\n", + "input_identifiers = \"\" #@param {type:\"string\"}\n", + "#@markdown Enter the identifiers (comma-separated), run the cell then you can:\n", + "#@markdown - Enter the path to a directory containing the images for the identifier\n", + "#@markdown - Leave the input field empty, then press ENTER to choose the file upload option\n", + "\n", + "\n", + "output_folder = \"renamed_instance_images\"\n", + "print(f\"Output folder: ${output_folder}\")\n", + "os.makedirs(output_folder, exist_ok=True)\n", + "\n", + "\n", + "# Define a function to display the form for entering the image folder paths\n", + "def display_input_form(identifiers_list):\n", + " identifier_paths_dict = {}\n", + " for identifier in identifiers_list:\n", + " path = input(f\"Path, or leave empty for file upload for identifier '{identifier}'': \").strip()\n", + " identifier_paths_dict[identifier] = path\n", + " \n", + " return identifier_paths_dict\n", + " \n", + "def rename_files(identifier, input_files, output_folder):\n", + " for i, file_name in enumerate(input_files):\n", + " if not os.path.isfile(file_name):\n", + " continue\n", + " _, ext = os.path.splitext(file_name)\n", + " new_file_name = f'{identifier} ({i + 1}){ext}'\n", + " output_file_name = os.path.join(output_folder, new_file_name)\n", + " copyfile(file_name, output_file_name)\n", + " print(f'Copied {file_name} -> {output_file_name}')\n", + "\n", + "\n", + "if root_folder_path:\n", + " identifiers = [folder for folder in os.listdir(root_folder_path) if os.path.isdir(os.path.join(root_folder_path, folder))]\n", + " \n", + " for identifier in identifiers:\n", + " path = os.path.join(root_folder_path, identifier, '*')\n", + " all_files = glob(path)\n", + " print(all_files)\n", + " all_files = [file_path for file_path in all_files if os.path.isfile(file_path)]\n", + " rename_files(identifier, all_files, output_folder)\n", + "else:\n", + " identifiers_list = [i.strip() for i in input_identifiers.split(',')]\n", + " paths_dict = display_input_form(identifiers_list)\n", + " for identifier in identifiers_list:\n", + " path = paths_dict[identifier]\n", + " if path == '':\n", + " print(f\"Upload the instance images for {identifier}: \")\n", + " uploaded = files.upload()\n", + " uploaded_files = list(uploaded.keys())\n", + " rename_files(identifier, uploaded_files, output_folder)\n", + " elif path:\n", + " all_files = glob(os.path.join(path, '*'))\n", + " all_files = [file_path for file_path in all_files if os.path.isfile(file_path)]\n", + " rename_files(identifier, all_files, output_folder)\n", + " else:\n", + " print(f\"Invalid input for identifier {identifier}, skipping.\")\n", + "\n", + "#@markdown ----\n", + "#@markdown - You can add `external captions` in txt files by simply giving each txt file the same name as the instance image, for example dikgur (1).jpg and dikgur (1).txt, and upload them here, to use the external captions, check the box \"external_captions\" in the training cell. `All the images must have one same extension` jpg or png or....etc\n", "\n", - "#@markdown - If you prefer to specify directly the folder of the pictures instead of uploading, this will add the pictures to the existing (if any) instance images. Leave EMPTY to upload.\n", + "IMAGES_FOLDER_OPTIONAL= output_folder\n", "\n", "Smart_Crop_images= True #@param{type: 'boolean'}\n", "Crop_size = 512 #@param [\"512\", \"576\", \"640\", \"704\", \"768\", \"832\", \"896\", \"960\", \"1024\"] {type:\"raw\"}\n", @@ -1616,10 +1689,6 @@ "metadata": { "accelerator": "GPU", "colab": { - "collapsed_sections": [ - "bbKbx185zqlz", - "AaLtXBbPleBr" - ], "provenance": [] }, "kernelspec": {