{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "61e8ae81-7b3f-435f-80fe-c536dd2e4352",
   "metadata": {},
   "source": [
    "# Color space conversion\n",
    "\n",
    "* The terminology and the way equations are constructed are detailed in [the documentation](https://cutcutcodec.readthedocs.io/latest/build/api/cutcutcodec.core.colorspace.html#module-cutcutcodec.core.colorspace).\n",
    "* The available **gammas** and **gamuts** are also detailed in [the documentation](https://cutcutcodec.readthedocs.io/latest/build/api/cutcutcodec.core.colorspace.cst.html). They come from International Telecommunication Union reports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e7491441-4e33-439b-8ecd-0f075e9460f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Math\n",
    "\n",
    "import numpy as np\n",
    "import sympy\n",
    "import torch\n",
    "\n",
    "import cutcutcodec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "172249d4-63bd-4061-bf85-be0868549597",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Source Colorspace: Colorspace(\"y'pbpr\", 'ntsc', 'smpte240m')\n",
      "Target Colorspace: Colorspace(\"y'pbpr\", 'bt2020', 'bt1361e, bt1361')\n"
     ]
    }
   ],
   "source": [
    "src = cutcutcodec.Colorspace(\"y'pbpr\", \"smpte240m\", \"smpte240m\")\n",
    "dst = cutcutcodec.Colorspace(\"y'pbpr\", \"bt2020\", \"bt2020\")\n",
    "print(f\"Source Colorspace: {src}\")\n",
    "print(f\"Target Colorspace: {dst}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0024c53-a356-4ff8-b810-a29c222341e8",
   "metadata": {},
   "source": [
    "## Symbolic conversions\n",
    "\n",
    "This allows you to retrieve the passage equations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "843203a0-c1f4-43e3-a066-569f91e9430d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace(\"y'pbpr\", 'ntsc', 'smpte240m') -> Colorspace(\"r'g'b'\", 'ntsc', 'smpte240m')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\frac{430238494 p_{r}}{273127803} + y', \\  - \\frac{7623978355136 p_{b}}{33806856032607} - \\frac{209727006492056 p_{r}}{439489128423891} + y', \\  \\frac{38383246 p_{b}}{21009831} + y'\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace(\"r'g'b'\", 'ntsc', 'smpte240m') -> Colorspace('rgb', 'ntsc')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\begin{cases} 0.25 r' & \\text{for}\\: r' \\leq 0.09128634211778008882172239358977671184 \\\\0.790526217176035190857986080494319491 \\left(r' + 0.111572195921731219670994036609727092\\right)^{2.222222222222222222222222222222222222} & \\text{otherwise} \\end{cases}, \\  \\begin{cases} 0.25 g' & \\text{for}\\: g' \\leq 0.09128634211778008882172239358977671184 \\\\0.790526217176035190857986080494319491 \\left(g' + 0.111572195921731219670994036609727092\\right)^{2.222222222222222222222222222222222222} & \\text{otherwise} \\end{cases}, \\  \\begin{cases} 0.25 b' & \\text{for}\\: b' \\leq 0.09128634211778008882172239358977671184 \\\\0.790526217176035190857986080494319491 \\left(b' + 0.111572195921731219670994036609727092\\right)^{2.222222222222222222222222222222222222} & \\text{otherwise} \\end{cases}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace('rgb', 'ntsc') -> Colorspace('xyz')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\frac{4026032 b}{21009831} + \\frac{99764014 g}{273127803} + \\frac{35828814 r}{91042601}, \\  \\frac{1818208 b}{21009831} + \\frac{191482543 g}{273127803} + \\frac{58008556 r}{273127803}, \\  \\frac{20130160 b}{21009831} + \\frac{30572843 g}{273127803} + \\frac{1706134 r}{91042601}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace('xyz') -> Colorspace('rgb', 'bt2020')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\frac{21532150939 x}{12543355000} - \\frac{4461219061 y}{12543355000} - \\frac{3178002061 z}{12543355000}, \\  - \\frac{1976779337 x}{2965129750} + \\frac{4793012913 y}{2965129750} + \\frac{46755163 z}{2965129750}, \\  \\frac{79263327 x}{4492356500} - \\frac{192186423 y}{4492356500} + \\frac{4233267077 z}{4492356500}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace('rgb', 'bt2020') -> Colorspace(\"r'g'b'\", 'bt2020', 'bt1361e, bt1361')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\begin{cases} 4.5 r & \\text{for}\\: r \\leq 0.01805396851080780733586959258468773494 \\\\1.099296826809442940347282759215782542 r^{0.45} - 0.099296826809442940347282759215782542 & \\text{otherwise} \\end{cases}, \\  \\begin{cases} 4.5 g & \\text{for}\\: g \\leq 0.01805396851080780733586959258468773494 \\\\1.099296826809442940347282759215782542 g^{0.45} - 0.099296826809442940347282759215782542 & \\text{otherwise} \\end{cases}, \\  \\begin{cases} 4.5 b & \\text{for}\\: b \\leq 0.01805396851080780733586959258468773494 \\\\1.099296826809442940347282759215782542 b^{0.45} - 0.099296826809442940347282759215782542 & \\text{otherwise} \\end{cases}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace(\"r'g'b'\", 'bt2020', 'bt1361e, bt1361') -> Colorspace(\"y'pbpr\", 'bt2020', 'bt1361e, bt1361')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\frac{826593596 b'}{13942086899} + \\frac{9452833643 g'}{13942086899} + \\frac{3662659660 r'}{13942086899}, \\  \\frac{b'}{2} - \\frac{859348513 g'}{2384635146} - \\frac{166484530 r'}{1192317573}, \\  - \\frac{413296798 b'}{10279427239} - \\frac{9452833643 g'}{20558854478} + \\frac{r'}{2}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "steps = [  # details of each transition stage\n",
    "    src,\n",
    "    cutcutcodec.Colorspace(\"r'g'b'\", src.primaries, src.transfer),\n",
    "    cutcutcodec.Colorspace(\"rgb\", src.primaries),\n",
    "    cutcutcodec.Colorspace(\"xyz\"),\n",
    "    cutcutcodec.Colorspace(\"rgb\", dst.primaries),\n",
    "    cutcutcodec.Colorspace(\"r'g'b'\", dst.primaries, dst.transfer),\n",
    "    dst,\n",
    "]\n",
    "for step_n, step_np1 in zip(steps[:-1], steps[1:]):\n",
    "    print(f\"{step_n} -> {step_np1}\")\n",
    "    eq = step_n.to_equation(step_np1)\n",
    "    display(Math(sympy.latex(eq)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0e4adfa-0851-47e4-b18a-f2fd9c035f0e",
   "metadata": {},
   "source": [
    "## Numerical conversion\n",
    "\n",
    "* Offers more flexibility than the FilterVideoColorspace.\n",
    "* Compiled in C, this function optimizes cache and overhead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f66adfe1-fbb7-4236-82aa-d1c6c0c0dc29",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def lambdify(u, v, y):\n",
      "    # this section is not cached and compiled in C\n",
      "    _0 = 0.25*y\n",
      "    _1 = 1.575227747868641553126687728674769884*v\n",
      "    _2 = _1 + y\n",
      "    _2 = _2 <= 0.09128634211778008882172239358977671184\n",
      "    _3 = 0.6348288375144786556453500415980909416*y\n",
      "    _4 = _3 + v + 0.07082924743593028657554055134201062229\n",
      "    _4 = _4**2.222222222222222222222222222222222222\n",
      "    _1 = 0.393806936967160388281671932168692471*v\n",
      "    _5 = _0 + _1\n",
      "    _6 = 2.169986653760008383303681644539441443*_4\n",
      "    _5 = Piecewise((_5, _2), (_6, True))\n",
      "    _3 = -_0\n",
      "    _7 = -y\n",
      "    _1 = 0.4772063583101252952258757092198383568*v\n",
      "    _8 = 0.2255157459120897865523043115836921751*u\n",
      "    _9 = _1 + _7 + _8\n",
      "    _10 = _9 >= -0.09128634211778008882172239358977671184\n",
      "    _9 = -_9\n",
      "    _9 = _9 + 0.111572195921731219670994036609727092\n",
      "    _9 = _9**2.222222222222222222222222222222222222\n",
      "    _7 = -_3\n",
      "    _1 = -0.1193015895775313238064689273049595892*v\n",
      "    _8 = -0.05637893647802244663807607789592304378*u\n",
      "    _11 = _1 + _7 + _8\n",
      "    _12 = 0.790526217176035190857986080494319491*_9\n",
      "    _11 = Piecewise((_11, _10), (_12, True))\n",
      "    _8 = 1.826918360266677061800259126310916066*u\n",
      "    _13 = _8 + y\n",
      "    _13 = _13 <= 0.09128634211778008882172239358977671184\n",
      "    _7 = 0.547369834224025763740773774057566679*y\n",
      "    _14 = _7 + u + 0.06107125438568854110753062185460237414\n",
      "    _14 = _14**2.222222222222222222222222222222222222\n",
      "    _8 = 0.4567295900666692654500647815777290165*u\n",
      "    _15 = _0 + _8\n",
      "    _16 = 3.016576666019477949253855712670004261*_14\n",
      "    _15 = Piecewise((_15, _13), (_16, True))\n",
      "    _12 = _10 & _13 & _2\n",
      "    _8 = 0.005616235079712466335963477384737015402*u\n",
      "    _1 = 0.1927478152272564917376197799690692823*v\n",
      "    _17 = _10 & _13\n",
      "    _6 = 1.291728221795239713514927736752038444*_4\n",
      "    _18 = -0.04167368848927251101449620675513716215*v\n",
      "    _19 = _18 + _6\n",
      "    _20 = _13 & _2\n",
      "    _21 = 0.02531017415011096215842454741477759892*u\n",
      "    _22 = 0.2761416962997584172930209810800754179*_9\n",
      "    _18 = 0.2344215037165290027521159867242064444*v\n",
      "    _23 = _18 + _22\n",
      "    _22 = _22 + _6\n",
      "    _24 = _10 & _2\n",
      "    _16 = 0.1671669241814816596465504740706525702*_14\n",
      "    _25 = -0.01969393907039849582246107003004058352*u\n",
      "    _26 = _16 + _25\n",
      "    _6 = 2.678715552469682602446063496570518666*_5\n",
      "    _27 = 1.571911982613223125966637466400477675*_11\n",
      "    _28 = 0.2493724649170942715872990370290036597*_15\n",
      "    _27 = _27 + _28 + _6\n",
      "    _29 = _0 + _1 + _8\n",
      "    _29 = _29 <= 0.01805396851080780733586959258468773494\n",
      "    _7 = 0.1011824693072398554196631390794156297*y\n",
      "    _30 = _19 + _7 + _8\n",
      "    _30 = _30 <= 0.01805396851080780733586959258468773494\n",
      "    _7 = 0.1626715565214876041129645851999734625*y\n",
      "    _31 = _21 + _23 + _7\n",
      "    _31 = _31 <= 0.01805396851080780733586959258468773494\n",
      "    _7 = 0.01385402582872745953262772427938909221*y\n",
      "    _32 = _21 + _22 + _7\n",
      "    _32 = _32 <= 0.01805396851080780733586959258468773494\n",
      "    _7 = 0.2361459741712725404673722757206109078*y\n",
      "    _33 = _1 + _26 + _7\n",
      "    _33 = _33 <= 0.01805396851080780733586959258468773494\n",
      "    _7 = 0.08732844347851239588703541480002653748*y\n",
      "    _34 = _19 + _26 + _7\n",
      "    _34 = _34 <= 0.01805396851080780733586959258468773494\n",
      "    _7 = 0.1488175306927601445803368609205843703*y\n",
      "    _23 = _16 + _23 + _7\n",
      "    _23 = _23 <= 0.01805396851080780733586959258468773494\n",
      "    _22 = _16 + _22\n",
      "    _22 = _22 <= 0.01805396851080780733586959258468773494\n",
      "    _22 = ITE(_2, _23, _22)\n",
      "    _22 = ITE(_10, _34, _22)\n",
      "    _22 = ITE(_24, _33, _22)\n",
      "    _22 = ITE(_13, _32, _22)\n",
      "    _22 = ITE(_20, _31, _22)\n",
      "    _22 = ITE(_17, _30, _22)\n",
      "    _22 = ITE(_12, _29, _22)\n",
      "    _16 = 0.09309404452711730793041398121145267568*_15\n",
      "    _23 = 0.5868155658274261080830700630726390165*_11\n",
      "    _23 = _16 + _23 + _5\n",
      "    _23 = _23**0.45\n",
      "    _23 = 0.8704349342486382027632689798713363333*_23\n",
      "    _23 = _23 - 0.099296826809442940347282759215782542\n",
      "    _22 = Piecewise((_27, _22), (_23, True))\n",
      "    _1 = 0.07436298549573965448441799964934607902*v\n",
      "    _8 = 0.03781848215812902175715342131049288056*u\n",
      "    _21 = 0.0124438037631045686968790177159063483*u\n",
      "    _23 = 0.7047606293783494066015658260648247585*_9\n",
      "    _18 = 0.03199536357565577701992240690343537089*v\n",
      "    _27 = _18 + _23\n",
      "    _18 = 0.1063583490713954315043404065527814499*v\n",
      "    _19 = 0.1763034254197555500027145811118515679*_4\n",
      "    _23 = _19 + _23\n",
      "    _25 = 0.05026228592123359045403243902639922886*u\n",
      "    _16 = 0.08218799238040872356530390459586689964*_14\n",
      "    _29 = 4.011787039184756895396630033053938588*_11\n",
      "    _6 = 0.3656084303625597013891725274218072362*_5\n",
      "    _26 = 0.1226045304526834032141974395242541763*_15\n",
      "    _29 = _26 + _29 + _6\n",
      "    _30 = _1 + _3 + _8\n",
      "    _30 = _30 >= -0.01805396851080780733586959258468773494\n",
      "    _7 = 0.0271229422675135058112983314970034118*y\n",
      "    _31 = _21 + _27 + _7\n",
      "    _31 = _31 <= 0.01805396851080780733586959258468773494\n",
      "    _6 = -_19\n",
      "    _7 = -0.2296884205354133499228237484765662647*y\n",
      "    _32 = _18 + _6 + _7 + _8\n",
      "    _32 = _32 >= -0.01805396851080780733586959258468773494\n",
      "    _7 = 0.006811362802926855734122079973569676459*y\n",
      "    _33 = _21 + _23 + _7\n",
      "    _33 = _33 <= 0.01805396851080780733586959258468773494\n",
      "    _26 = -_16\n",
      "    _7 = -0.2431886371970731442658779200264303235*y\n",
      "    _34 = _1 + _25 + _26 + _7\n",
      "    _34 = _34 >= -0.01805396851080780733586959258468773494\n",
      "    _7 = 0.02031157946458665007717625152343373534*y\n",
      "    _27 = _16 + _27 + _7\n",
      "    _27 = _27 <= 0.01805396851080780733586959258468773494\n",
      "    _1 = -_18\n",
      "    _8 = -_25\n",
      "    _7 = 0.2228770577324864941887016685029965882*y\n",
      "    _35 = _1 + _16 + _19 + _7 + _8\n",
      "    _35 = _35 <= 0.01805396851080780733586959258468773494\n",
      "    _23 = _16 + _23\n",
      "    _23 = _23 <= 0.01805396851080780733586959258468773494\n",
      "    _35 = ITE(_10, _35, _23)\n",
      "    _35 = ITE(_2, _27, _35)\n",
      "    _35 = ITE(_24, _34, _35)\n",
      "    _35 = ITE(_13, _33, _35)\n",
      "    _35 = ITE(_17, _32, _35)\n",
      "    _35 = ITE(_20, _31, _35)\n",
      "    _35 = ITE(_12, _30, _35)\n",
      "    _19 = 0.09113355888323915291044422946884672446*_5\n",
      "    _16 = 0.0305610764617251742427671942086476572*_15\n",
      "    _23 = _11 + _16 + _19\n",
      "    _23 = _23**0.45\n",
      "    _23 = 1.043930016251889889020578059853752167*_23\n",
      "    _23 = _23 - 0.099296826809442940347282759215782542\n",
      "    _35 = Piecewise((_29, _35), (_23, True))\n",
      "    _8 = 0.4076023144074993873444640152111481931*u\n",
      "    _1 = 0.003664093591595216860355383337237092768*v\n",
      "    _21 = 0.004619273747905038815867694158849998107*u\n",
      "    _14 = 2.722613229148498087964592165258763969*_14\n",
      "    _18 = 0.006110597724791278555250393601747546687*v\n",
      "    _9 = 0.0647698809191872617237235775360352408*_9\n",
      "    _25 = 0.4122215881554044261603317093699981912*u\n",
      "    _23 = _25 + _9\n",
      "    _9 = _14 + _9\n",
      "    _19 = 0.03367110699321960713282169771998237656*_4\n",
      "    _36 = -0.009774691316386495415605776938984639455*v\n",
      "    _4 = _19 + _36\n",
      "    _16 = 4.061477922699390175209513706070768804*_15\n",
      "    _6 = 0.06982530570266157964103038662075216118*_5\n",
      "    _27 = 0.3686967715979482451494559073084790346*_11\n",
      "    _27 = _16 + _27 + _6\n",
      "    _36 = -_1\n",
      "    _29 = _0 + _36 + _8\n",
      "    _29 = _29 <= 0.01805396851080780733586959258468773494\n",
      "    _16 = -_14\n",
      "    _0 = -0.02436233762781165693280479410717951087*y\n",
      "    _30 = _0 + _1 + _16 + _21\n",
      "    _30 = _30 >= -0.01805396851080780733586959258468773494\n",
      "    _0 = 0.2295168460223362086028080051495289425*y\n",
      "    _31 = _0 + _18 + _23\n",
      "    _31 = _31 <= 0.01805396851080780733586959258468773494\n",
      "    _0 = 0.003879183650147865535612799256708453399*y\n",
      "    _32 = _0 + _18 + _9\n",
      "    _32 = _32 <= 0.01805396851080780733586959258468773494\n",
      "    _0 = 0.2461208163498521344643872007432915466*y\n",
      "    _33 = _0 + _4 + _8\n",
      "    _33 = _33 <= 0.01805396851080780733586959258468773494\n",
      "    _8 = -_21\n",
      "    _0 = 0.02048315397766379139719199485047105748*y\n",
      "    _34 = _0 + _14 + _4 + _8\n",
      "    _34 = _34 <= 0.01805396851080780733586959258468773494\n",
      "    _0 = 0.2256376623721883430671952058928204891*y\n",
      "    _23 = _0 + _19 + _23\n",
      "    _23 = _23 <= 0.01805396851080780733586959258468773494\n",
      "    _9 = _19 + _9\n",
      "    _9 = _9 <= 0.01805396851080780733586959258468773494\n",
      "    _9 = ITE(_13, _23, _9)\n",
      "    _9 = ITE(_10, _34, _9)\n",
      "    _9 = ITE(_17, _33, _9)\n",
      "    _9 = ITE(_2, _32, _9)\n",
      "    _9 = ITE(_20, _31, _9)\n",
      "    _9 = ITE(_24, _30, _9)\n",
      "    _9 = ITE(_12, _29, _9)\n",
      "    _10 = 0.09077896731564661384689587977436820083*_11\n",
      "    _2 = 0.01719209288628938625683614354268095072*_5\n",
      "    _10 = _10 + _15 + _2\n",
      "    _10 = _10**0.45\n",
      "    _10 = 1.049728979382482414626981540495914549*_10\n",
      "    _10 = _10 - 0.099296826809442940347282759215782542\n",
      "    _9 = Piecewise((_27, _9), (_10, True))\n",
      "    _10 = 0.2627052669039600712074144417509988725*_22\n",
      "    _11 = 0.6780070811119393525736767106632850431*_35\n",
      "    _12 = 0.05928765198410057621890884758571608441*_9\n",
      "    _10 = _10 + _11 + _12\n",
      "    _11 = 0.5*_9\n",
      "    _12 = -0.1396310293247686621154916082076398282*_22\n",
      "    _17 = -0.3603689706752313378845083917923601718*_35\n",
      "    _11 = _11 + _12 + _17\n",
      "    _12 = 0.5*_22\n",
      "    _35 = -0.4597937911917934633089336856193296705*_35\n",
      "    _9 = -0.04020620880820653669106631438067032949*_9\n",
      "    _35 = _12 + _35 + _9\n",
      "    _ = (_10, _11, _35)\n",
      "    return _\n"
     ]
    }
   ],
   "source": [
    "func = src.to_function(dst)  # optimize the full chain\n",
    "print(func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bc4e62b5-99b5-4e36-8277-dcaaabb2bf89",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_yuv = torch.rand(2160, 3840, 3)\n",
    "src_yuv[..., 1:] -= 0.5\n",
    "src_yuv *= 0.2  # valid range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b15c50c4-821a-4eae-b5ab-c92548c44a85",
   "metadata": {},
   "outputs": [],
   "source": [
    "dst_y, dst_u, dst_v = func(y=src_yuv[..., 0], u=src_yuv[..., 1], v=src_yuv[..., 2])\n",
    "dst_yuv = torch.cat([dst_y[..., None], dst_u[..., None], dst_v[..., None]], dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5caba756-82c1-4173-8326-06cecf08d730",
   "metadata": {},
   "source": [
    "## Alternatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "93ad2ab6-d4ab-4325-bed5-1de1e2d117e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import timeit\n",
    "\n",
    "import colour  # pip install colour-science"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6e628f5b-a15a-40f3-ad67-b1d275732bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def colourscience_convert(src_yuv):\n",
    "    src_colourspace = colour.models.RGB_COLOURSPACE_SMPTE_240M\n",
    "    src_k = colour.WEIGHTS_YCBCR[\"SMPTE-240M\"]\n",
    "    dst_colourspace = colour.models.RGB_COLOURSPACE_BT2020\n",
    "    dst_k = colour.WEIGHTS_YCBCR[\"ITU-R BT.2020\"]\n",
    "    src_rgb = colour.YCbCr_to_RGB(src_yuv, src_k, in_range=(0.0, 1.0, -0.5, 0.5), out_range=(0.0, 1.0))\n",
    "    xyz = colour.RGB_to_XYZ(src_rgb, src_colourspace, apply_cctf_decoding=True)\n",
    "    dst_rgb = colour.XYZ_to_RGB(xyz, dst_colourspace, apply_cctf_encoding=True)\n",
    "    dst_yuv = colour.RGB_to_YCbCr(dst_rgb, dst_k, in_range=(0.0, 1.0), out_range=(0.0, 1.0, -0.5, 0.5))\n",
    "    return dst_yuv\n",
    "\n",
    "dst_yuv_ = colourscience_convert(src_yuv.numpy(force=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "423fab4c-4f8e-4e0d-99f7-e72254bee40d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(3.3795e-05, dtype=torch.float64)\n",
      "tensor(0.0001, dtype=torch.float64)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_614229/2375015244.py:1: DeprecationWarning: __array_wrap__ must accept context and return_scalar arguments (positionally) in the future. (Deprecated NumPy 2.0)\n",
      "  print(abs(dst_yuv[..., 2] - dst_yuv_[..., 2]).mean())\n",
      "/tmp/ipykernel_614229/2375015244.py:2: DeprecationWarning: __array_wrap__ must accept context and return_scalar arguments (positionally) in the future. (Deprecated NumPy 2.0)\n",
      "  print(abs(dst_yuv[..., 2] - dst_yuv_[..., 2]).max())\n"
     ]
    }
   ],
   "source": [
    "print(abs(dst_yuv[..., 2] - dst_yuv_[..., 2]).mean())\n",
    "print(abs(dst_yuv[..., 2] - dst_yuv_[..., 2]).max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b2f3627a-65b2-4cea-864b-8d724ce88254",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cutcutcodec convesion take 280.15 ms\n",
      "colourscience convesion take 7838.57 ms\n",
      "cutcutcodec is 27.97981168338215 times faster than colourscience\n"
     ]
    }
   ],
   "source": [
    "number = 5\n",
    "time_cutcutcodec = timeit.repeat(lambda: func(y=src_yuv[..., 0], u=src_yuv[..., 1], v=src_yuv[..., 2]), number=number, repeat=7)\n",
    "print(f\"cutcutcodec convesion take {1000*np.median(time_cutcutcodec)/number:.2f} ms\")\n",
    "time_colourscience = timeit.repeat(lambda: colourscience_convert(src_yuv.numpy(force=True)), number=number, repeat=5)\n",
    "print(f\"colourscience convesion take {1000*np.median(time_colourscience)/number:.2f} ms\")\n",
    "print(f\"cutcutcodec is {np.median(time_colourscience)/np.median(time_cutcutcodec)} times faster than colourscience\")"
   ]
  }
 ],
 "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.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}