{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "23717cf2",
   "metadata": {},
   "source": [
    "# 🐱 Catrex Lite — удаление фона\n",
    "\n",
    "Демо модели [Catniti/catrex-lite-image-segmentation](https://huggingface.co/Catniti/catrex-lite-image-segmentation).\n",
    "\n",
    "Компактная сегментационная сеть (4.63M параметров), обученная **с нуля** на\n",
    "DIS5K. Вырезает объект из фона и отдаёт прозрачный PNG.\n",
    "\n",
    "**Два способа запуска:**\n",
    "\n",
    "| | ONNX | PyTorch |\n",
    "|---|---|---|\n",
    "| Зависимости | `onnxruntime` (~50 МБ) | `torch` (~2.5 ГБ) |\n",
    "| GPU | не нужен | не нужен |\n",
    "| Скорость на CPU | ~0.4 с | ~0.6 с |\n",
    "\n",
    "Ниже сначала ONNX — он проще и легче. PyTorch-вариант в конце, если нужно\n",
    "дообучать или встраивать в свой пайплайн.\n",
    "\n",
    "> Работает и без GPU. `Среда выполнения` → `Выполнить все`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f04a4ba",
   "metadata": {},
   "source": [
    "## Установка"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0e9fd44",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q onnxruntime pillow numpy huggingface_hub\n",
    "print(\"готово\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d88d917",
   "metadata": {},
   "source": [
    "## Способ 1: ONNX (рекомендуется)\n",
    "\n",
    "Модель экспортирована с динамическими осями, поэтому принимает любой размер\n",
    "входа. Но обучалась она на 512×512 — на этом разрешении результат лучше всего,\n",
    "поэтому приводим к нему, а маску потом растягиваем обратно."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cec250b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import onnxruntime as ort\n",
    "from PIL import Image\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "REPO_ID = \"Catniti/catrex-lite-image-segmentation\"\n",
    "\n",
    "model_path = hf_hub_download(REPO_ID, \"model.onnx\")\n",
    "session = ort.InferenceSession(model_path, providers=[\"CPUExecutionProvider\"])\n",
    "\n",
    "IMG_SIZE = 512\n",
    "MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)\n",
    "STD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)\n",
    "\n",
    "\n",
    "def remove_background(image, threshold=None, feather=True):\n",
    "    '''Возвращает (RGBA без фона, маска). threshold=None -> мягкая альфа.'''\n",
    "    img = image.convert(\"RGB\")\n",
    "\n",
    "    x = np.asarray(img.resize((IMG_SIZE, IMG_SIZE), Image.BILINEAR), dtype=np.float32) / 255.0\n",
    "    x = ((x - MEAN) / STD).transpose(2, 0, 1)[None]\n",
    "\n",
    "    mask = session.run(None, {\"input\": x.astype(np.float32)})[0][0, 0]\n",
    "\n",
    "    # маску считаем в размере модели, затем растягиваем — так тонкие детали\n",
    "    # сохраняются лучше, чем если гонять через сеть уменьшенную картинку\n",
    "    mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(img.size, Image.BILINEAR)\n",
    "    mask_np = np.asarray(mask)\n",
    "\n",
    "    if threshold is not None:\n",
    "        mask_np = np.where(mask_np > threshold * 255, 255, 0).astype(np.uint8)\n",
    "        mask = Image.fromarray(mask_np)\n",
    "\n",
    "    cutout = Image.fromarray(np.dstack([np.asarray(img), mask_np]), \"RGBA\")\n",
    "    return cutout, mask\n",
    "\n",
    "\n",
    "print(\"модель загружена:\", model_path.split(\"/\")[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58503b78",
   "metadata": {},
   "source": [
    "## Пробуем на примере"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b6e0921",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io, urllib.request\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "URL = \"https://huggingface.co/datasets/mishig/sample_images/resolve/main/tiger.jpg\"\n",
    "sample = Image.open(io.BytesIO(urllib.request.urlopen(URL).read()))\n",
    "\n",
    "cutout, mask = remove_background(sample)\n",
    "\n",
    "def show_on_checkerboard(rgba, size=16):\n",
    "    '''Шахматка под прозрачностью — иначе на белом фоне ничего не видно.'''\n",
    "    w, h = rgba.size\n",
    "    tile = np.indices((h, w)).sum(axis=0) // size % 2\n",
    "    bg = np.where(tile[..., None], 205, 255).astype(np.uint8).repeat(3, axis=2)\n",
    "    bg = Image.fromarray(bg)\n",
    "    bg.paste(rgba, (0, 0), rgba)\n",
    "    return bg\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(15, 5))\n",
    "ax[0].imshow(sample);                      ax[0].set_title(\"оригинал\")\n",
    "ax[1].imshow(mask, cmap=\"gray\");           ax[1].set_title(\"маска\")\n",
    "ax[2].imshow(show_on_checkerboard(cutout));ax[2].set_title(\"фон удалён\")\n",
    "for a in ax: a.axis(\"off\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ab14910",
   "metadata": {},
   "source": [
    "## Своё изображение\n",
    "\n",
    "Запусти ячейку и выбери файл. Результат скачается автоматически."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad632958",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from google.colab import files\n",
    "    uploaded = files.upload()\n",
    "except ImportError:\n",
    "    uploaded = {}\n",
    "    print(\"не Colab — подставь свой путь: Image.open('photo.jpg')\")\n",
    "\n",
    "for name in uploaded:\n",
    "    img = Image.open(io.BytesIO(uploaded[name]))\n",
    "    cutout, mask = remove_background(img)\n",
    "\n",
    "    out_name = f\"nobg_{name.rsplit('.', 1)[0]}.png\"\n",
    "    cutout.save(out_name)\n",
    "\n",
    "    fig, ax = plt.subplots(1, 3, figsize=(15, 5))\n",
    "    ax[0].imshow(img.convert(\"RGB\"));           ax[0].set_title(\"оригинал\")\n",
    "    ax[1].imshow(mask, cmap=\"gray\");            ax[1].set_title(\"маска\")\n",
    "    ax[2].imshow(show_on_checkerboard(cutout)); ax[2].set_title(\"без фона\")\n",
    "    for a in ax: a.axis(\"off\")\n",
    "    plt.tight_layout(); plt.show()\n",
    "\n",
    "    files.download(out_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "889a1ac1",
   "metadata": {},
   "source": [
    "## Замена фона\n",
    "\n",
    "Раз есть альфа-канал, объект можно положить на что угодно."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e77592bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "cutout, _ = remove_background(sample)\n",
    "\n",
    "variants = {\n",
    "    \"белый\":  Image.new(\"RGB\", cutout.size, (255, 255, 255)),\n",
    "    \"чёрный\": Image.new(\"RGB\", cutout.size, (18, 18, 18)),\n",
    "    \"цвет\":   Image.new(\"RGB\", cutout.size, (99, 102, 241)),\n",
    "}\n",
    "\n",
    "fig, ax = plt.subplots(1, len(variants), figsize=(15, 5))\n",
    "for a, (title, bg) in zip(ax, variants.items()):\n",
    "    composed = bg.copy()\n",
    "    composed.paste(cutout, (0, 0), cutout)\n",
    "    a.imshow(composed); a.set_title(title); a.axis(\"off\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "536ebb63",
   "metadata": {},
   "source": [
    "## Способ 2: PyTorch\n",
    "\n",
    "Нужен, если планируешь дообучать модель или встраивать её в существующий\n",
    "torch-пайплайн. Архитектура объявлена прямо здесь — отдельного пакета\n",
    "устанавливать не надо."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c9c1ecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q torch safetensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d04433fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class ConvBNReLU(nn.Module):\n",
    "    def __init__(self, cin, cout, dilation=1):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Conv2d(cin, cout, 3, padding=dilation, dilation=dilation)\n",
    "        self.bn   = nn.BatchNorm2d(cout)\n",
    "    def forward(self, x):\n",
    "        return F.relu(self.bn(self.conv(x)), inplace=True)\n",
    "\n",
    "def _up(x, ref):\n",
    "    return F.interpolate(x, size=ref.shape[2:], mode=\"bilinear\", align_corners=False)\n",
    "\n",
    "class RSU(nn.Module):\n",
    "    '''Residual U-block: маленький U-Net внутри слоя.'''\n",
    "    def __init__(self, depth, cin, cmid, cout):\n",
    "        super().__init__()\n",
    "        self.depth = depth\n",
    "        self.head  = ConvBNReLU(cin, cout)\n",
    "        self.enc   = nn.ModuleList([ConvBNReLU(cout if i == 0 else cmid, cmid)\n",
    "                                    for i in range(depth)])\n",
    "        self.bottom = ConvBNReLU(cmid, cmid, dilation=2)\n",
    "        self.dec = nn.ModuleList([ConvBNReLU(cmid * 2, cmid if i > 0 else cout)\n",
    "                                  for i in range(depth)])\n",
    "        self.pool = nn.MaxPool2d(2, 2, ceil_mode=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.head(x)\n",
    "        skips = []\n",
    "        h = x\n",
    "        for i, e in enumerate(self.enc):\n",
    "            h = e(h)\n",
    "            skips.append(h)\n",
    "            if i < self.depth - 1:\n",
    "                h = self.pool(h)\n",
    "        h = self.bottom(h)\n",
    "        for i in range(self.depth - 1, -1, -1):\n",
    "            h = self.dec[i](torch.cat([h, skips[i]], 1))\n",
    "            if i > 0:\n",
    "                h = _up(h, skips[i - 1])\n",
    "        return h + x\n",
    "\n",
    "class CatrexLite(nn.Module):\n",
    "    def __init__(self, ch=(16, 32, 64, 128, 256)):\n",
    "        super().__init__()\n",
    "        c1, c2, c3, c4, c5 = ch\n",
    "        self.pool = nn.MaxPool2d(2, 2, ceil_mode=True)\n",
    "\n",
    "        self.e1 = RSU(5, 3,  c1 // 2, c1)\n",
    "        self.e2 = RSU(4, c1, c1 // 2, c2)\n",
    "        self.e3 = RSU(3, c2, c2 // 2, c3)\n",
    "        self.e4 = RSU(3, c3, c3 // 2, c4)\n",
    "        self.e5 = RSU(2, c4, c4 // 2, c5)\n",
    "        self.e6 = RSU(2, c5, c5 // 2, c5)\n",
    "\n",
    "        self.d5 = RSU(2, c5 * 2, c4 // 2, c4)\n",
    "        self.d4 = RSU(3, c4 * 2, c3 // 2, c3)\n",
    "        self.d3 = RSU(3, c3 * 2, c2 // 2, c2)\n",
    "        self.d2 = RSU(4, c2 * 2, c1 // 2, c1)\n",
    "        self.d1 = RSU(5, c1 * 2, c1 // 2, c1)\n",
    "\n",
    "        self.side = nn.ModuleList([\n",
    "            nn.Conv2d(c, 1, 3, padding=1) for c in (c1, c1, c2, c3, c4, c5)])\n",
    "        self.fuse = nn.Conv2d(6, 1, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h1 = self.e1(x)\n",
    "        h2 = self.e2(self.pool(h1))\n",
    "        h3 = self.e3(self.pool(h2))\n",
    "        h4 = self.e4(self.pool(h3))\n",
    "        h5 = self.e5(self.pool(h4))\n",
    "        h6 = self.e6(self.pool(h5))\n",
    "\n",
    "        u5 = self.d5(torch.cat([_up(h6, h5), h5], 1))\n",
    "        u4 = self.d4(torch.cat([_up(u5, h4), h4], 1))\n",
    "        u3 = self.d3(torch.cat([_up(u4, h3), h3], 1))\n",
    "        u2 = self.d2(torch.cat([_up(u3, h2), h2], 1))\n",
    "        u1 = self.d1(torch.cat([_up(u2, h1), h1], 1))\n",
    "\n",
    "        feats = [u1, u2, u3, u4, u5, h6]\n",
    "        sides = [_up(s(f), x) for s, f in zip(self.side, feats)]\n",
    "        return [self.fuse(torch.cat(sides, 1))] + sides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42b9870b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from safetensors.torch import load_file\n",
    "\n",
    "cfg = json.load(open(hf_hub_download(REPO_ID, \"config.json\")))\n",
    "weights = load_file(hf_hub_download(REPO_ID, \"model.safetensors\"))\n",
    "\n",
    "torch_model = CatrexLite(ch=tuple(cfg[\"channels\"]))\n",
    "torch_model.load_state_dict(weights)\n",
    "torch_model.eval()\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "torch_model.to(device)\n",
    "\n",
    "print(f\"параметров: {sum(p.numel() for p in torch_model.parameters())/1e6:.2f}M | {device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f78669f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def remove_background_torch(image):\n",
    "    img = image.convert(\"RGB\")\n",
    "    x = np.asarray(img.resize((IMG_SIZE, IMG_SIZE), Image.BILINEAR), dtype=np.float32) / 255.0\n",
    "    x = torch.from_numpy(((x - MEAN) / STD).transpose(2, 0, 1)[None]).to(device)\n",
    "\n",
    "    # сеть отдаёт 7 выходов (deep supervision), нужен только первый\n",
    "    mask = torch.sigmoid(torch_model(x)[0])[0, 0].cpu().numpy()\n",
    "\n",
    "    mask = Image.fromarray((mask * 255).astype(np.uint8)).resize(img.size, Image.BILINEAR)\n",
    "    return Image.fromarray(np.dstack([np.asarray(img), np.asarray(mask)]), \"RGBA\"), mask\n",
    "\n",
    "\n",
    "cutout_t, mask_t = remove_background_torch(sample)\n",
    "\n",
    "# сверяем с ONNX — расхождение должно быть на уровне погрешности\n",
    "diff = np.abs(np.asarray(mask_t, dtype=np.float32) - np.asarray(mask, dtype=np.float32)).mean()\n",
    "print(f\"среднее расхождение ONNX vs PyTorch: {diff:.4f} (из 255)\")\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(10, 5))\n",
    "ax[0].imshow(mask_t, cmap=\"gray\");            ax[0].set_title(\"маска (PyTorch)\")\n",
    "ax[1].imshow(show_on_checkerboard(cutout_t)); ax[1].set_title(\"без фона\")\n",
    "for a in ax: a.axis(\"off\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b742429",
   "metadata": {},
   "source": [
    "## Пакетная обработка\n",
    "\n",
    "Для папки с изображениями."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "508534d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "def process_folder(src_dir, dst_dir=\"output\"):\n",
    "    src, dst = Path(src_dir), Path(dst_dir)\n",
    "    dst.mkdir(exist_ok=True)\n",
    "\n",
    "    exts = {\".jpg\", \".jpeg\", \".png\", \".webp\", \".bmp\"}\n",
    "    files_list = [p for p in sorted(src.iterdir()) if p.suffix.lower() in exts]\n",
    "\n",
    "    for i, p in enumerate(files_list, 1):\n",
    "        cutout, _ = remove_background(Image.open(p))\n",
    "        cutout.save(dst / f\"{p.stem}.png\")\n",
    "        print(f\"[{i}/{len(files_list)}] {p.name}\")\n",
    "\n",
    "    print(f\"готово -> {dst}/\")\n",
    "\n",
    "# process_folder(\"my_images\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe0c3a4f",
   "metadata": {},
   "source": [
    "## Ограничения\n",
    "\n",
    "Модель обучалась на одной T4 в разрешении 512px, F1 на DIS-VD = **0.6461**.\n",
    "\n",
    "Где работает хорошо:\n",
    "- чёткий одиночный объект на контрастном фоне\n",
    "- товарные фото, предметы, техника\n",
    "\n",
    "Где будет слабее:\n",
    "- волосы, мех, перья — тонкие структуры размываются\n",
    "- прозрачные и полупрозрачные объекты\n",
    "- объект сливается с фоном по цвету\n",
    "\n",
    "Если нужно максимальное качество — посмотри\n",
    "[BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet) (MIT) или\n",
    "[RMBG-2.0](https://huggingface.co/briaai/RMBG-2.0) (CC BY-NC, только\n",
    "некоммерческое использование). Они обучались неделями на A100/H200.\n",
    "\n",
    "---\n",
    "\n",
    "Обучающий ноутбук и код: [страница модели](https://huggingface.co/Catniti/catrex-lite-image-segmentation)."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}