Upload pixai_tagger.ipynb
Browse files
colab_notebooks/utils/pixai_tagger.ipynb
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"authorship_tag":"ABX9TyNIMUUXZrdmas+n6CIimymY"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"code","execution_count":null,"metadata":{"id":"T_aHXn_z9tE3"},"outputs":[],"source":["#@title PixAI Tagger Batch Processor for Google Drive Zip\n","#@markdown Mount Google Drive and process images with pixai-tagger-v0.9\n","\n","# Input / Output paths (Google Drive)\n","input_zip_path = \"/content/drive/MyDrive/input_images.zip\" #@param {type:\"string\"}\n","output_zip_path = \"/content/drive/MyDrive/tagged_output.zip\" #@param {type:\"string\"}\n","\n","# Optional: thresholds\n","general_threshold = 0.30 #@param {type:\"number\"}\n","character_threshold = 0.85 #@param {type:\"number\"}\n","\n","# ============================================================\n","# 1. Install dependencies & setup\n","# ============================================================\n","!pip install -q 'dghs-imgutils>=0.19.0' torch huggingface_hub timm pillow pandas onnxruntime-gpu 2>/dev/null || \\\n"," !pip install -q 'dghs-imgutils>=0.19.0' torch huggingface_hub timm pillow pandas onnxruntime\n","\n","from google.colab import drive, userdata\n","import os\n","import zipfile\n","import shutil\n","from pathlib import Path\n","from PIL import Image\n","from tqdm.auto import tqdm\n","from imgutils.tagging import get_pixai_tags\n","from huggingface_hub import login, hf_hub_download\n","import tempfile\n","\n","# Mount Drive\n","drive.mount('/content/drive', force_remount=False)\n","\n","# HF Token (prefer Colab Secrets)\n","try:\n"," HF_TOKEN = userdata.get('HF_TOKEN')\n","except Exception:\n"," HF_TOKEN = os.environ.get('HF_TOKEN') or os.environ.get('HUGGINGFACE_HUB_TOKEN')\n","\n","if HF_TOKEN:\n"," login(token=HF_TOKEN)\n"," os.environ['HF_TOKEN'] = HF_TOKEN\n"," os.environ['HUGGINGFACE_HUB_TOKEN'] = HF_TOKEN\n"," print(\"Hugging Face token set – faster downloads enabled.\")\n","else:\n"," print(\"WARNING: No HF_TOKEN found. Downloads may be slower / rate-limited.\")\n","\n","# ============================================================\n","# 2. Work directories\n","# ============================================================\n","work_dir = Path(\"/content/pixai_work\")\n","extract_dir = work_dir / \"extracted\"\n","output_dir = work_dir / \"output\"\n","\n","if work_dir.exists():\n"," shutil.rmtree(work_dir)\n","extract_dir.mkdir(parents=True)\n","output_dir.mkdir(parents=True)\n","\n","# ============================================================\n","# 3. Extract input zip\n","# ============================================================\n","print(f\"Extracting: {input_zip_path}\")\n","with zipfile.ZipFile(input_zip_path, 'r') as zf:\n"," zf.extractall(extract_dir)\n","\n","# Collect image files (common formats)\n","image_exts = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif', '.tiff', '.tif'}\n","image_files = []\n","for root, _, files in os.walk(extract_dir):\n"," for f in files:\n"," if Path(f).suffix.lower() in image_exts:\n"," image_files.append(Path(root) / f)\n","\n","image_files = sorted(image_files) # deterministic order\n","print(f\"Found {len(image_files)} images.\")\n","\n","if len(image_files) == 0:\n"," raise RuntimeError(\"No images found in the zip!\")\n","\n","# ============================================================\n","# 4. Process each image\n","# ============================================================\n","print(\"Loading PixAI Tagger (first run downloads the model)...\")\n","\n","for idx, img_path in enumerate(tqdm(image_files, desc=\"Tagging\")):\n"," num = idx + 1\n"," stem = img_path.stem # original basename without extension\n","\n"," # --- Load & convert image to RGB JPG ---\n"," try:\n"," img = Image.open(img_path).convert(\"RGB\")\n"," except Exception as e:\n"," print(f\"Skipping {img_path.name}: {e}\")\n"," continue\n","\n"," out_img_path = output_dir / f\"{num}.jpg\"\n"," img.save(out_img_path, quality=95, optimize=True)\n","\n"," # --- Get tags ---\n"," try:\n"," result = get_pixai_tags(\n"," img,\n"," model_name='v0.9',\n"," thresholds={'general': general_threshold, 'character': character_threshold},\n"," fmt=('general', 'character', 'ips')\n"," )\n"," # result can be tuple depending on fmt\n"," if isinstance(result, tuple):\n"," general = result[0] if len(result) > 0 else {}\n"," character = result[1] if len(result) > 1 else {}\n"," ips = result[2] if len(result) > 2 else []\n"," else:\n"," general, character, ips = {}, {}, []\n","\n"," # Collect tags (higher confidence first already sorted by the library)\n"," tag_list = []\n"," for t in general.keys():\n"," tag_list.append(t)\n"," for t in character.keys():\n"," tag_list.append(t)\n"," for t in ips:\n"," if t not in tag_list:\n"," tag_list.append(t)\n","\n"," # Format: \"tag1, tag2, tag3\" (space after every comma)\n"," new_tags = \", \".join(tag_list)\n","\n"," except Exception as e:\n"," print(f\"Tagging failed for {img_path.name}: {e}\")\n"," new_tags = \"\"\n","\n"," # --- Handle existing text file (same stem) ---\n"," existing_txt = None\n"," for candidate in [\n"," img_path.with_suffix('.txt'),\n"," img_path.parent / f\"{stem}.txt\",\n"," ]:\n"," if candidate.exists():\n"," existing_txt = candidate\n"," break\n","\n"," # Also search the whole extract tree for a matching .txt\n"," if existing_txt is None:\n"," for root, _, files in os.walk(extract_dir):\n"," for f in files:\n"," if f.lower() == f\"{stem}.txt\".lower():\n"," existing_txt = Path(root) / f\n"," break\n"," if existing_txt:\n"," break\n","\n"," final_text = new_tags\n"," if existing_txt and existing_txt.exists():\n"," with open(existing_txt, 'r', encoding='utf-8', errors='ignore') as f:\n"," old_text = f.read().strip()\n"," if old_text:\n"," # Append new tags, keep space after commas\n"," if old_text.endswith(','):\n"," final_text = old_text + \" \" + new_tags\n"," else:\n"," final_text = old_text + \", \" + new_tags if new_tags else old_text\n","\n"," # Ensure consistent spacing after commas\n"," final_text = \", \".join([t.strip() for t in final_text.split(\",\") if t.strip()])\n","\n"," # Write numbered text file\n"," out_txt_path = output_dir / f\"{num}.txt\"\n"," with open(out_txt_path, 'w', encoding='utf-8') as f:\n"," f.write(final_text)\n","\n","print(f\"Processed {len(list(output_dir.glob('*.jpg')))} images.\")\n","\n","# ============================================================\n","# 5. Create output zip on Drive\n","# ============================================================\n","print(f\"Creating output zip: {output_zip_path}\")\n","os.makedirs(os.path.dirname(output_zip_path), exist_ok=True)\n","\n","with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:\n"," for file in sorted(output_dir.iterdir()):\n"," zf.write(file, arcname=file.name)\n","\n","print(\"Done!\")\n","print(f\"Output saved to: {output_zip_path}\")\n","print(f\"Contents: numbered 1.jpg / 1.txt , 2.jpg / 2.txt , ...\")"]}]}
|