Torch Painter.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "torch.random.manual_seed(0)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from colormath.color_objects import sRGBColor, LabColor\n",
    "from colormath.color_conversions import convert_color\n",
    "from colormath.color_diff import delta_e_cie2000\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.cluster import KMeans\n",
    "import cv2\n",
    "import os\n",
    "import pandas as pd\n",
    "import copy\n",
    "\n",
    "\n",
    "\n",
    "def discrete_color(color_stroke, allowed_colors, just_inds=False): #(n*5, 3)\n",
    "    ''' color_stroke in RGB\n",
    "    allowed_colors in BGR '''\n",
    "    allowed_colors_tensors = [np.array([allowed_colors[i]] * color_stroke.shape[0]) for i in range(len(allowed_colors))]\n",
    "    \n",
    "    l2_distances = np.zeros((color_stroke.shape[0], len(allowed_colors_tensors)))\n",
    "    for i in range(len(allowed_colors_tensors)):\n",
    "        l2_distances[:, i] = np.sum((color_stroke - allowed_colors_tensors[i])**2, axis=1)\n",
    "        for j in range(l2_distances.shape[0]):\n",
    "            color1_rgb = sRGBColor(color_stroke[j,0], color_stroke[j,1], color_stroke[j,2])\n",
    "            color2_rgb = sRGBColor(allowed_colors[i][2], allowed_colors[i][1], allowed_colors[i][0])\n",
    "            color1_lab = convert_color(color1_rgb, LabColor)\n",
    "            color2_lab = convert_color(color2_rgb, LabColor)\n",
    "            l2_distances[j, i] = delta_e_cie2000(color1_lab, color2_lab)\n",
    "\n",
    "    color_inds = np.tile(np.argmin(l2_distances, axis=1)[np.newaxis].T,(1,3))\n",
    "    if just_inds:\n",
    "        return color_inds\n",
    "    \n",
    "    new_color_stroke = np.zeros(color_stroke.shape)\n",
    "    for i in range(len(allowed_colors_tensors)):\n",
    "        new_color_stroke = np.where(color_inds == i, allowed_colors_tensors[i], new_color_stroke)\n",
    "\n",
    "    return new_color_stroke\n",
    "\n",
    "\n",
    "def color_cluster(img, n_colors=10):\n",
    "    \"\"\"\n",
    "    Cluster the colors of a given image\n",
    "    \n",
    "    args:\n",
    "        img (np.array[width, height] : image in BGR\n",
    "    kwargs:\n",
    "        n_colors (int) : number of colors to cluster\n",
    "        \n",
    "    return:\n",
    "        List((B,G,R),...) : List of BGR tuples\n",
    "    \"\"\"\n",
    "    allowed_colors = []\n",
    "\n",
    "    colors = img.reshape((img.shape[0]*img.shape[1]), 3) / 255.\n",
    "\n",
    "    kmeans = KMeans(n_clusters=n_colors)\n",
    "    kmeans.fit(colors)\n",
    "\n",
    "    for i in range(n_colors):\n",
    "        c = kmeans.cluster_centers_[i] # c is in BGR format\n",
    "        allowed_colors.append(c) #BGR format appended\n",
    "    return allowed_colors # They're global anyways\n",
    "\n",
    "def save_colors(allowed_colors, actions, output_dir):\n",
    "    \"\"\"\n",
    "    Save the colors used as an image so you know how to mix the paints\n",
    "    \n",
    "    args:\n",
    "        allowed_colors (List((B,G,R),...) : List of BGR tuples\n",
    "        actions (List(actions)) : list of 13 valued tuples. Used to determine how much each color is used\n",
    "        output_dir (String) : Where to save the image\n",
    "    \"\"\"\n",
    "    n_colors = len(allowed_colors)\n",
    "    fig, ax = plt.subplots(2, 6, figsize=(1*12, 5))\n",
    "    i = 0 \n",
    "    w = 128\n",
    "    for c in allowed_colors:\n",
    "        # print('[', int(255*c[2]), ', ', int(255*c[1]), ', ', int(255*c[0]), '],', end='', sep='')\n",
    "        num_uses = np.sum(actions[:,12] == i)\n",
    "        ax[i//6, i%6].imshow(np.concatenate((np.ones((w,w,1))*c[2], np.ones((w,w,1))*c[1], np.ones((w,w,1))*c[0]), axis=-1))\n",
    "        ax[i//6, i%6].set_xticks([])\n",
    "        ax[i//6, i%6].set_yticks([])\n",
    "        ax[i//6, i%6].set_title(i)\n",
    "        ax[i//6, i%6].set_xlabel(str(num_uses) + ' uses')\n",
    "        i += 1\n",
    "    while i < 12:\n",
    "        ax[i//6, i%6].imshow(np.concatenate((np.ones((w,w,1)), np.ones((w,w,1)), np.ones((w,w,1))), axis=-1))\n",
    "        ax[i//6, i%6].set_xticks([])\n",
    "        ax[i//6, i%6].set_yticks([])\n",
    "        ax[i//6, i%6].set_title('Not used')\n",
    "        i += 1\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.savefig(os.path.join(output_dir, 'colors.png'))\n",
    "\n",
    "def group_strokes_by_color(strokes):\n",
    "    ''' Sort the strokes by color '''\n",
    "    strokes = sorted(copy.deepcopy(strokes),key=lambda l:l[12], reverse=False)\n",
    "    return strokes\n",
    "\n",
    "MAX_ITERATIONS = 100\n",
    "def draw(f, width=128, height=128):\n",
    "    \"\"\"\n",
    "    Draw a stroke onto a blank, 1-channel canvas\n",
    "    \n",
    "    args:\n",
    "        f (Tuple) : x0, y0, x1, y1, x2, y2, z0, z2, w0, w2\n",
    "        r (int) : radius in pixels of stroke\n",
    "    \n",
    "    kwargs:\n",
    "        width (int) :  Width of canvas. (Default 128)\n",
    "        height (int) :  Height of canvas. (Default 128)\n",
    "    \n",
    "    return:\n",
    "        np.array[width, height] : matrix (boolean map) with the stroke drawn on it.\n",
    "    \"\"\"\n",
    "    x0, y0, x1, y1, x2, y2, z0, z2, w0, w2 = f\n",
    "    x0, y0, x1, y1, x2, y2 = x0*width, y0*height, x1*width, y1*height, x2*width, y2*height\n",
    "    canvas = np.zeros([width, height]).astype('float32')\n",
    "    for i in range(MAX_ITERATIONS):\n",
    "        t = i/MAX_ITERATIONS\n",
    "        x = (int)((1-t) * (1-t) * x0 + 2 * t * (1-t) * x1 + t * t * x2)\n",
    "        y = (int)((1-t) * (1-t) * y0 + 2 * t * (1-t) * y1 + t * t * y2)\n",
    "        z = max(int(z0 / 1.65), 1)\n",
    "        w = 1.\n",
    "        cv2.circle(canvas, (y, x), z, w, -1)\n",
    "    return 1 - canvas\n",
    "\n",
    "def paint_discrete_colors(S, allowed_colors, img):\n",
    "    \"\"\"\n",
    "    Paint an image using only the allowed _colors\n",
    "    \n",
    "    args:\n",
    "        S (List[Tuple 13]) : List of strokes. x0,y0,x1,y1,x2,y2,r0,r1,opaq0,opaq1,B,G,R\n",
    "        allowed_colors (List((B,G,R),...) : List of BGR tuples\n",
    "        img (np.array[width, height, 3]) : target image, used for shape\n",
    "    \n",
    "    return:\n",
    "        np.array[width, height] : matrix (boolean map) with the stroke drawn on it.\n",
    "    \"\"\"\n",
    "    canvas = torch.ones(img.shape)\n",
    "    S = copy.deepcopy(S)\n",
    "    S[:,8:10] = 1. # opaque\n",
    "    for K in S:\n",
    "        K[10:13] = allowed_colors[int(K[11])]\n",
    "        s = 1 - draw(K[:10], width=img.shape[0], height=img.shape[1])\n",
    "        canvas = apply_stroke(canvas, torch.from_numpy(s), torch.from_numpy(K[10:13].astype(float)))\n",
    "    return canvas\n",
    "\n",
    "def save_strokes_for_robot(S, img, painting, allowed_colors, ignore_whites=True, output_dir='peters_paint_instructions'):\n",
    "    \"\"\"\n",
    "    Save the strokes, colors, painting, and target image for the robot\n",
    "    \n",
    "    args:        \n",
    "        img (np.array[width, height, 3]) : target image RGB 0-255\n",
    "        S (List[Tuple 13]) : List of strokes. x0,y0,x1,y1,x2,y2,r0,r1,opaq0,opaq1,B,G,R\n",
    "        painting (np.array[width, height, 3]) : Painting RGB 0-1\n",
    "        allowed_colors (List((B,G,R),...) : List of BGR tuples\n",
    "    \n",
    "    kwargs:\n",
    "        ignore_whites (bool) :  Ignore strokes that are pretty much white\n",
    "        output_dir (str) : where to save all these files\n",
    "    \"\"\"\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.mkdir(output_dir)\n",
    "    \n",
    "    # Remove white strokes if specified\n",
    "    if ignore_whites:\n",
    "        S_no_white = []\n",
    "        S_whites = np.array(copy.deepcopy(S))\n",
    "        S_whites[:,10:13] = discrete_color(S_whites[:,10:13] / 255., allowed_colors, just_inds=False)\n",
    "        for i in range(len(S_whites)):\n",
    "            if S_whites[i,10:13].mean() < 0.95:\n",
    "                S_no_white.append(S[i])\n",
    "        S = S_no_white\n",
    "    S = np.array(copy.deepcopy(S))\n",
    "    \n",
    "    # Discritize the colors, replace BGR with paint color indices\n",
    "    color_inds = discrete_color(S[:,10:13] / 255., allowed_colors, just_inds=True)\n",
    "    S[:,10:13] = color_inds\n",
    "    \n",
    "    # Scale based on canvas size\n",
    "    S[:,[0,2,4]] /= img.shape[0]\n",
    "    S[:,[1,3,5]] /= img.shape[1]\n",
    "    S[:,0:6] = np.clip(S[:,0:6], a_min=0, a_max=1)\n",
    "    S = np.nan_to_num(S)\n",
    "\n",
    "    cv2.imwrite(os.path.join(output_dir, 'painting.jpg'), cv2.resize(painting[:,:,::-1] * 255., (img.shape[0]*4, img.shape[1]*4)))\n",
    "    cv2.imwrite(os.path.join(output_dir, 'target.jpg'), cv2.resize(img[:,:,::-1], (img.shape[0]*4, img.shape[1]*4)))\n",
    "    discrete_painting = paint_discrete_colors(S, allowed_colors, img)*255.\n",
    "    cv2.imwrite(os.path.join(output_dir, 'discrete_painting.jpg'), cv2.resize(discrete_painting.cpu().numpy(), (img.shape[0]*4, img.shape[1]*4)))\n",
    "\n",
    "    save_colors(allowed_colors, S, output_dir)\n",
    "    plt.show()\n",
    "    df = pd.DataFrame(S)\n",
    "    # Sort by brush, then color\n",
    "    df = df.sort_values([6, 12], ascending=[False, True])\n",
    "\n",
    "    df.to_csv(os.path.join(output_dir, 'actions.csv'), sep=\",\", header=False, index=False, float_format='%.5f')\n",
    "\n",
    "def discretize_image(img, allowed_colors):\n",
    "    \"\"\"\n",
    "    Only use allowed_colors in the given image. Use euclidean distance for speed.\n",
    "    \n",
    "    args:\n",
    "        img (np.array[width, height, 3]) : target image BGR\n",
    "        allowed_colors (List((B,G,R),...) : List of BGR tuples\n",
    "    \n",
    "    return:\n",
    "        np.array[width, height, 3] : RGB 0-255 np.uint8 target image using only the allowed colors\n",
    "    \"\"\"\n",
    "    n_pix = img.shape[0]*img.shape[1]\n",
    "    n_colors = len(allowed_colors)\n",
    "    \n",
    "    img_flat = np.reshape(img, (n_pix, 3)) / 255.\n",
    "\n",
    "    color_mat = np.empty((n_colors, n_pix, 3))\n",
    "    \n",
    "    i = 0\n",
    "    for c in allowed_colors:\n",
    "        color_mat[i] = np.tile(c[np.newaxis].T, (1, n_pix)).T\n",
    "        i += 1\n",
    "\n",
    "    img_exp = np.tile(img_flat[np.newaxis], (n_colors, 1, 1))\n",
    "    diff = np.sum(np.abs(img_exp - color_mat), axis=2)\n",
    "\n",
    "    argmin = np.argmin(diff, axis=0)\n",
    "\n",
    "    img_disc = np.array(allowed_colors)[argmin]\n",
    "    img_disc = np.reshape(img_disc, (img.shape[0],img.shape[1], 3))\n",
    "#     plt.imshow(img_disc[:,:,::-1])\n",
    "#     plt.show()\n",
    "    \n",
    "    return (img_disc[:,:,::-1] * 255.).astype(np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.ndimage.filters import gaussian_filter\n",
    "from sklearn.decomposition import PCA\n",
    "import warnings\n",
    "warnings.simplefilter('ignore', np.RankWarning)\n",
    "np.seterr(divide='ignore', invalid='ignore')\n",
    "\n",
    "def draw_spline_stroke(K, r, width=128, height=128):\n",
    "    \"\"\"\n",
    "    Paint a stroke defined by a list of points onto a canvas\n",
    "    \n",
    "    args:\n",
    "        K (List[Tup(int, int)]) : a nested list of points to draw. [(x_pixel, y_pixel),...]\n",
    "        r (int) : radius in pixels of stroke\n",
    "    \n",
    "    kwargs:\n",
    "        width (int) :  Width of canvas. (Default 128)\n",
    "        height (int) :  Height of canvas. (Default 128)\n",
    "    \n",
    "    return:\n",
    "        np.array[width, height] : matrix (boolean map) with the stroke drawn on it.\n",
    "    \"\"\"\n",
    "    canvas = np.zeros([width, height]).astype('float32')\n",
    "\n",
    "    for f in K:\n",
    "        x = f[0]\n",
    "        y = f[1]\n",
    "        z = max(int(r / 1.65), 1)\n",
    "        w = 1.\n",
    "        cv2.circle(canvas, (y, x), z, w, -1)\n",
    "    return torch.from_numpy(1 - canvas).to(device)\n",
    "\n",
    "def make_stroke(x0, y0, r, ref_image, canvas, max_stroke_length=None):\n",
    "    \"\"\"\n",
    "    Draw a curved line on a canvas from a starting point based on how it decreases loss\n",
    "    \n",
    "    args:\n",
    "        x0 (int) : Starting x pixel\n",
    "        y1 (int) : Starting y pixel\n",
    "        r (int) : Radius in pixels of stroke\n",
    "        ref_image (np.array[width, height, 3]) :  Reference image 0-255 RGB\n",
    "        canvas (np.array[width, height, 3]) :  Current painting canvas 0-1 RGB\n",
    "    \n",
    "    kwargs:\n",
    "        max_stroke_length (int) : Maximum length of a stroke in pixels.\n",
    "    \n",
    "    return:\n",
    "        [(x0,y0),..] : List of x,y coordinates in the stroke\n",
    "    \"\"\"\n",
    "    stroke_color = ref_image[x0,y0,:]\n",
    "    K = [(x0,y0)]\n",
    "    \n",
    "    x, y = x0, y0\n",
    "    \n",
    "    # default max stroke length is 1/3rd of canvas width\n",
    "    max_stroke_length = max_stroke_length if max_stroke_length is not None else int(ref_image.shape[0] * 0.15)\n",
    "    min_stroke_length = int(ref_image.shape[0] * 0.05)\n",
    "    \n",
    "    width, height, _ = ref_image.shape\n",
    "    \n",
    "    canvas_hat = canvas.clone() # Canvas with proposed stroke\n",
    "    original_loss = torch.mean(torch.abs(ref_image - canvas*255.))\n",
    "    prev_loss = original_loss.clone()\n",
    "    \n",
    "    grid = max(3*r, 4)\n",
    "    grid = grid + 1 if (grid%2+1)==0 else grid\n",
    "    \n",
    "    for i in range(1, max_stroke_length):\n",
    "        x = max(min(x, ref_image.shape[0]-1), 0)\n",
    "        y = max(min(y, ref_image.shape[1]-1), 0)\n",
    "        \n",
    "        # Determine where the stroke color will help canvas become reference image\n",
    "        D = ref_image[max(x-grid//2, 0):x+grid//2, max(y-grid//2, 0):y+grid//2] \\\n",
    "            - stroke_color[None, None, :]\n",
    "\n",
    "        # Multiply by diff of canvas and ref image so not painting over parts already painted\n",
    "        D = torch.abs(D) * (1 - torch.abs(ref_image[max(x-grid//2, 0):x+grid//2, max(y-grid//2, 0):y+grid//2] \\\n",
    "                 - canvas_hat[max(x-grid//2, 0):x+grid//2, max(y-grid//2, 0):y+grid//2])/(255.*3))\n",
    "        \n",
    "#         plt.matshow(torch.sum(torch.abs(D), dim=2).cpu())\n",
    "#         plt.colorbar()\n",
    "#         plt.show()\n",
    "        \n",
    "        # # Blur it doesn't seem to help\n",
    "        # D = torch.nn.functional.avg_pool2d(D, max(r//2, 2), stride=1)\n",
    "        \n",
    "        D = torch.sum(torch.abs(D), dim=2)\n",
    "        \n",
    "        D = D.float()\n",
    "        D[D.shape[0]//2, D.shape[1]//2] = 1000. #don't take from center\n",
    "        \n",
    "        noise = torch.randn(D.shape[0], D.shape[1], device=device)*0.0001\n",
    "        D = D + noise\n",
    "#         plt.matshow(D.cpu())\n",
    "#         plt.colorbar()\n",
    "#         plt.show()\n",
    "        \n",
    "        while(True): # Keep trying coordinates until you get one you haven't been to before\n",
    "            dx, dy = np.unravel_index(D.argmin().cpu(), D.shape)\n",
    "            og_dx, og_dy = dx, dy\n",
    "            dx -= D.shape[0]//2\n",
    "            dy -= D.shape[1]//2\n",
    "            if dx == 0 and dy == 0:\n",
    "                break # it's at the edge anyways :(\n",
    "            hyp = (dx**2 + dy**2)**(.5)\n",
    "            ratio = max(r/1.5,1.) / hyp\n",
    "            dx, dy = dx * ratio, dy * ratio\n",
    "            x_hat = int(max(x + dx, 0))\n",
    "            y_hat = int(max(y + dy, 0))\n",
    "\n",
    "            if ((x_hat,y_hat) in K):\n",
    "                D[og_dx, og_dy] = 1000. # Don't pick this point again\n",
    "            else:\n",
    "                break\n",
    "        \n",
    "        # Paint the stroke\n",
    "        s_hat = 1 - draw_spline_stroke(K + [(x_hat, y_hat)], r, width=width, height=height)\n",
    "        canvas_hat = apply_stroke(canvas, s_hat, stroke_color / 255.)\n",
    "        \n",
    "        loss = torch.mean(torch.abs(ref_image - canvas_hat*255.))\n",
    "        if (loss > prev_loss) and (i > min_stroke_length):\n",
    "            # New point doesn't decrease the loss, terminate\n",
    "            break\n",
    "        prev_loss = loss\n",
    "        \n",
    "        x, y = x_hat, y_hat\n",
    "        K.append((x,y))\n",
    "#     plt.imshow(canvas_hat.cpu())\n",
    "#     plt.show()\n",
    "    return K\n",
    "\n",
    "\n",
    "def apply_stroke(canvas, stroke, color):\n",
    "    \"\"\"\n",
    "    Apply a given stroke to the canvas with a given color\n",
    "    \n",
    "    args:\n",
    "        canvas (np.array[width, height, 3]) : Current painting canvas 0-1 RGB\n",
    "        stroke (np.array[width, height]) :  Stroke boolean map\n",
    "        color (np.array[3]) : RGB color to use for the brush stroke\n",
    "    \n",
    "    return:\n",
    "        np.array[width, height, 3] : Painting with additional stroke in 0-1 RGB format\n",
    "    \"\"\"\n",
    "    s_expanded = stroke.unsqueeze(2).repeat(1,1,3)\n",
    "    s_color = s_expanded * color[None, None, :]\n",
    "    return canvas * (1 - s_expanded) + s_color\n",
    "\n",
    "def curved_stroke_to_straight(K, color, r):\n",
    "    \"\"\"\n",
    "    Convert an arbitrary shape set of points to a straight line using PCA\n",
    "    \n",
    "    args:\n",
    "        K (List[Tup(int, int)]) : a nested list of points to draw. [(x_pixel, y_pixel),...]\n",
    "        color (np.array[3]) : BGR color to use for the brush stroke\n",
    "    \n",
    "    return:\n",
    "        List[Tup(int, int)] : Straightened version of input K. [(x_pixel, y_pixel),...]\n",
    "        Tuple : Stroke for the robot. x0,y0,x1,y1,x2,y2,r0,r1,opaq0,opaq1,B,G,R\n",
    "    \"\"\"\n",
    "    original_K = copy.deepcopy(K)\n",
    "    K = np.array(K).T\n",
    "    straight_K = []\n",
    "    \n",
    "    try:\n",
    "        # Project points onto best fitting line, then sort them to find start and end of line\n",
    "        pca = PCA(n_components=1, svd_solver='auto')\n",
    "        reduced = pca.fit_transform(K.T)\n",
    "\n",
    "        sorted_inds = np.argsort(reduced[:,0])\n",
    "\n",
    "        x0, y0 = K[0, sorted_inds[0]], K[1, sorted_inds[0]] \n",
    "        x1, y1 = K[0, sorted_inds[len(sorted_inds)//2]], K[1, sorted_inds[len(sorted_inds)//2]] \n",
    "        x2, y2 = K[0, sorted_inds[-1]], K[1, sorted_inds[-1]]\n",
    "    except:\n",
    "        x0, y0 = K[0, 0], K[1, 0] \n",
    "        x1, y1 = K[0, K.shape[1]//2], K[1, K.shape[1]//2] \n",
    "        x2, y2 = K[0, -1], K[1, -1]\n",
    "    \n",
    "#     try:\n",
    "#         # Extend the line by r\n",
    "#         # m, b = np.polyfit(K[0], K[1], 1)\n",
    "#         if (x2-x0) == 0: pass\n",
    "#         m = (y2-y0) / (x2-x0)\n",
    "#         x0, y0 = x0 - (r/2 / (1+m**2)**(.5)), y0 - (r/2 / (1/m**2 + 1)**(.5))\n",
    "#         x2, y2 = x2 + (r/2 / (1+m**2)**(.5)), y2 + (r/2 / (1/m**2 + 1)**(.5))\n",
    "#     except:\n",
    "#         pass\n",
    "    \n",
    "    # orient right to left cuz robot sux at painting left to right :(\n",
    "    if x2 - x0 < 0:\n",
    "        x_temp, y_temp = x0, y0\n",
    "        x0, y0 = x2, y2\n",
    "        x2, y2 = x_temp, y_temp\n",
    "\n",
    "    robot_stroke = (x0, y0, x1, y1, x2, y2, r, r, 1, 1, color[0], color[1], color[2])\n",
    "    \n",
    "    MAX_ITERATIONS = 100 # How smooth\n",
    "    for i in range(MAX_ITERATIONS):\n",
    "        t = i * (1./MAX_ITERATIONS)\n",
    "        x = (int)((1-t) * (1-t) * x0 + 2 * t * (1-t) * x1 + t * t * x2)\n",
    "        y = (int)((1-t) * (1-t) * y0 + 2 * t * (1-t) * y1 + t * t * y2)\n",
    "        straight_K.append((x,y))\n",
    "    return straight_K, robot_stroke\n",
    "\n",
    "def paint_layer(canvas, reference_image, r, T, curved, pix_diff_thresh=30):\n",
    "    \"\"\"\n",
    "    Go through the pixels and paint a layer of strokes with a given radius\n",
    "    \n",
    "    args:\n",
    "        canvas (np.array[width, height, 3]) : Current painting canvas 0-1 RGB\n",
    "        reference_image (np.array[width, height, 3]) :  Reference image 0-255 RGB\n",
    "        r (int) : Brush radius to use\n",
    "        T (int) : Max attempts to generate strokes\n",
    "        curved (bool) : Whether to use curved or straight brush strokes\n",
    "    kwargs:\n",
    "        # smooth_loss (int) : How many bad strokes before giving up\n",
    "        pix_diff_thresh (int) : small value reduces likely hood of picking a dark color\n",
    "    \n",
    "    return:\n",
    "        np.array[width, height, 3] : Painting in 0-1 RGB format\n",
    "        List[Tuple 13] : List of strokes. x0,y0,x1,y1,x2,y2,r0,r1,opaq0,opaq1,B,G,R\n",
    "    \"\"\"\n",
    "    S = []\n",
    "            \n",
    "    width, height, _ = canvas.shape\n",
    "        \n",
    "    # Keep track of where you've already painted\n",
    "    already_painted = torch.zeros((width, height), dtype=torch.bool, device=device)\n",
    "    \n",
    "    d_losses = []\n",
    "    \n",
    "    for i in range(T):\n",
    "        # sum the error near (x,y)\n",
    "        D = torch.sum(torch.abs(canvas*255. - reference_image), dim=2)\n",
    "        \n",
    "        # D = gaussian_filter(D, sigma=max(r//4, 2)) # Blur it\n",
    "        \n",
    "        D = D * (1-already_painted.float()) # Don't paint same are twice\n",
    "        \n",
    "        # Pick starting point where error is largest\n",
    "        noise = torch.randn(D.shape[0], D.shape[1], device=device)*0.001\n",
    "        x, y = np.unravel_index((torch.clamp(D, 0, pix_diff_thresh) + noise).argmax().cpu(), D.shape)\n",
    "        \n",
    "        color = reference_image[x,y,:] / 255.\n",
    "\n",
    "        K = make_stroke(x, y, r, reference_image, canvas)\n",
    "        if not curved:\n",
    "            K, robot_stroke = curved_stroke_to_straight(K, color*255., r)\n",
    "        s = 1 - draw_spline_stroke(K, r, width=width, height=height)\n",
    "        \n",
    "        loss_before_stroke = torch.mean(torch.abs(reference_image - canvas*255.))\n",
    "        \n",
    "        canvas_hat = apply_stroke(canvas, s, color)\n",
    "        \n",
    "        loss_after_stroke = torch.mean(torch.abs(reference_image - canvas_hat*255.))\n",
    "        \n",
    "        stroke_loss = loss_after_stroke - loss_before_stroke\n",
    "        \n",
    "        # Reject strokes that increase loss\n",
    "        if stroke_loss > 0:\n",
    "            continue\n",
    "\n",
    "        # Accept the stroke\n",
    "        if not curved:\n",
    "            S.append(robot_stroke)\n",
    "        else:\n",
    "            S.append(K)\n",
    "        canvas = canvas_hat\n",
    "        already_painted = already_painted | (s==1)\n",
    "        \n",
    "#         if len(S) % 50 == 0:\n",
    "#             plt.imshow(canvas)\n",
    "#             plt.show()\n",
    "#     plt.imshow(canvas)\n",
    "#     plt.show()\n",
    "    print(loss_after_stroke)\n",
    "    return canvas, S\n",
    "\n",
    "def paint(source_image, R, T=[100], curved=True):\n",
    "    \"\"\"\n",
    "    Paint a given image\n",
    "    \n",
    "    args:\n",
    "        source_image (np.array[width, height, 3]) : Target image 0-255 RGB\n",
    "        R (list(int)) : List of brush radii to use\n",
    "    kwargs:\n",
    "        T (List[int]) : Maximum number of strokes per radius\n",
    "        curved (bool) : Whether to use curved or straight brush strokes\n",
    "    \n",
    "    return:\n",
    "        np.array[width, height, 3] : Painting in 0-1 RGB format\n",
    "    \"\"\"\n",
    "    canvas = torch.ones(source_image.shape).to(device)\n",
    "    \n",
    "    S = []\n",
    "\n",
    "    # paint the canvas\n",
    "    i = 0\n",
    "    for r in R:\n",
    "        # apply Gaussian blur, makes it more accurate, but if robot painting, you don't want to do it cuz colors\n",
    "#         reference_image = cv2.GaussianBlur(source_image, (r,r) if r%2 == 1 else (r+1, r+1), 0)\n",
    "#         reference_image = torch.from_numpy(reference_image).to(device)\n",
    "        reference_image = torch.from_numpy(source_image).to(device)\n",
    "        \n",
    "        # paint using the radius brush size\n",
    "        canvas, s = paint_layer(canvas, reference_image, r, T=T[i], curved=curved)\n",
    "        S += s\n",
    "        i += 1\n",
    "    return canvas, S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_frames(content_video_fn, animation_fps=10):\n",
    "    ''' Given video file name, return list of PIL images '''\n",
    "    video = cv2.VideoCapture(content_video_fn)\n",
    "\n",
    "    fps = video.get(cv2.CAP_PROP_FPS)\n",
    "\n",
    "    frame_skip = fps // animation_fps\n",
    "\n",
    "    frames, i = [], 0\n",
    "    success, image = video.read()\n",
    "\n",
    "    while success:\n",
    "        if i % frame_skip == 0:\n",
    "            frames.append(image)\n",
    "        i += 1\n",
    "        success, image = video.read()\n",
    "\n",
    "    return frames\n",
    "\n",
    "def resize_img(img, w=200, h=200):\n",
    "    ''' Reshape img BGR, remove alpha channel if exists '''\n",
    "    if img.shape[2] == 4:\n",
    "        trans_mask = img[:,:,3] == 0\n",
    "        img[trans_mask] = [255, 255, 255, 255]\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)\n",
    "    img = cv2.resize(img, (w,h))\n",
    "    return img\n",
    "\n",
    "def animate(video_fn, R, T=[100], n_colors=6, animation_fps=10, output_dir='animation_instructions'):\n",
    "    \"\"\"\n",
    "    Paint each frame from a given directory\n",
    "    \n",
    "    args:\n",
    "        video_fn : video file name\n",
    "        source_image (np.array[width, height, 3]) : Target image 0-255 RGB\n",
    "        R (list(int)) : List of brush radii to use\n",
    "    kwargs:\n",
    "        T (List[int]) : Maximum number of strokes per radius\n",
    "        curved (bool) : Whether to use curved or straight brush strokes\n",
    "    \n",
    "    return:\n",
    "        np.array[width, height, 3] : Painting in 0-1 RGB format\n",
    "    \"\"\"\n",
    "    frames = get_frames(video_fn, animation_fps=animation_fps)\n",
    "    \n",
    "    allowed_colors = color_cluster(frames[0], n_colors=n_colors)\n",
    "    allowed_colors = sorted(allowed_colors, key=lambda l:np.mean(l), reverse=False)\n",
    "    \n",
    "    if not os.path.exists(output_dir): os.mkdir(output_dir)\n",
    "    \n",
    "    i = 0\n",
    "    for img in frames:\n",
    "        img = resize_img(img)[:,:,::-1] # Resize and BGR to RBG\n",
    "        img = discretize_image(img[:,:,::-1], allowed_colors) # Convert pixels to closest allowed color\n",
    "        \n",
    "        painting, s = paint(img, R, T=T, curved=False)\n",
    "        \n",
    "        save_strokes_for_robot(s, img, painting.cpu().numpy(), allowed_colors,\n",
    "                               ignore_whites=True, output_dir=os.path.join(output_dir, str(i)))\n",
    "        i += 1\n",
    "# animate('atta1.mp4', [6], output_dir='animation_instructions_test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# animate('videos/atta_stylized.mp4', [6], T=[250], animation_fps=5, output_dir='animation_instructions_test')\n",
    "# animate('videos/atta1.mp4', [6], T=[250], animation_fps=5, output_dir='animation_instructions_test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "def paint_fn(fn):\n",
    "    img = cv2.imread(fn, cv2.IMREAD_UNCHANGED)\n",
    "    if img.shape[2] == 4:\n",
    "        trans_mask = img[:,:,3] == 0\n",
    "        img[trans_mask] = [255, 255, 255, 255]\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)\n",
    "    img = cv2.resize(img, (500,500))\n",
    "    img = img[:,:,::-1]\n",
    "    img = cv2.resize(img, (200,200))\n",
    "\n",
    "    fig, ax = plt.subplots(1,3, figsize=(10,20))\n",
    "    ax[0].imshow(cv2.resize(img, (500,500)))\n",
    "    ax[0].title.set_text('Original')\n",
    "\n",
    "    allowed_colors = color_cluster(img[:,:,::-1], n_colors=12)\n",
    "    # So when running it a few times, similar results\n",
    "    allowed_colors = sorted(allowed_colors, key=lambda l:np.mean(l), reverse=False)\n",
    "\n",
    "    img = discretize_image(img[:,:,::-1], allowed_colors)\n",
    "    # img[np.mean(img, axis=2) > 100, :] = 255\n",
    "    ax[1].imshow(cv2.resize(img, (500,500)))\n",
    "    ax[1].title.set_text('discrete')\n",
    "\n",
    "    painting, s = paint(img, [6], T=[1000], curved=False)\n",
    "\n",
    "    ax[2].imshow(cv2.resize(painting.cpu().numpy(), (500,500)))\n",
    "    ax[2].title.set_text('Painting' + str(len(s)))\n",
    "\n",
    "    [a.set_yticks([]) + a.set_xticks([]) for a in ax]\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    #save_strokes_for_robot(s, img, painting.cpu().numpy(), allowed_colors, ignore_whites=True, output_dir='nathan')\n",
    "\n",
    "for fn in ['../../HumanoidPainter/image/ye.jpg', '../../HumanoidPainter/image/vangogh.png', '../../Downloads/hertzmanns/images/tomato.jpg', \\\n",
    "           '../../HumanoidPainter/image/hambone.jpg', '../../Downloads/hertzmanns/Coolin_AnimationSequence/Assignment30016.png']:\n",
    "    paint_fn(fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "for fn in ['../../HumanoidPainter/image/ye.jpg', '../../HumanoidPainter/image/vangogh.png', '../../Downloads/hertzmanns/images/tomato.jpg', \\\n",
    "           '../../HumanoidPainter/image/hambone.jpg', '../../Downloads/hertzmanns/Coolin_AnimationSequence/Assignment30016.png']:\n",
    "    paint_fn(fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "if not os.path.exists('coolin_paint_instructions'):\n",
    "    os.mkdir('coolin_paint_instructions')\n",
    "i = 0\n",
    "\n",
    "\n",
    "img = cv2.imread('Coolin_AnimationSequence/Assignment30016.png', cv2.IMREAD_UNCHANGED)[80:-80,450:-150,:]#[:,:,::-1]\n",
    "if img.shape[2] == 4:\n",
    "    trans_mask = img[:,:,3] == 0\n",
    "    img[trans_mask] = [255, 255, 255, 255]\n",
    "    img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)\n",
    "img = cv2.resize(img, (500,500))\n",
    "img = img[:,:,::-1]\n",
    "allowed_colors = color_cluster(img[:,:,::-1], n_colors=6)\n",
    "# So when running it a few times, similar results\n",
    "allowed_colors = sorted(allowed_colors, key=lambda l:np.mean(l), reverse=False)\n",
    "    \n",
    "for fn in os.listdir('Coolin_AnimationSequence'):\n",
    "    print(i)\n",
    "    img = cv2.imread('Coolin_AnimationSequence/' + fn, cv2.IMREAD_UNCHANGED)[80:-80,450:-150,:]#[:,:,::-1]\n",
    "    \n",
    "    if img.shape[2] == 4:\n",
    "        trans_mask = img[:,:,3] == 0\n",
    "        img[trans_mask] = [255, 255, 255, 255]\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)\n",
    "    img = cv2.resize(img, (200,200))\n",
    "    img = img[:,:,::-1]\n",
    "    \n",
    "    fig, ax = plt.subplots(1,3)\n",
    "    ax[0].imshow(img)\n",
    "    ax[0].title.set_text('Original')\n",
    "\n",
    "    img = discretize_image(img[:,:,::-1], allowed_colors)\n",
    "\n",
    "    ax[1].imshow(img)\n",
    "    ax[1].title.set_text('discrete')\n",
    "\n",
    "    painting, s = paint(img, [6], T=[800], curved=False)\n",
    "\n",
    "    ax[2].imshow(painting.cpu())\n",
    "    ax[2].title.set_text('Painting' + str(len(s)))\n",
    "\n",
    "    [a.set_yticks([]) + a.set_xticks([]) for a in ax]\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    save_strokes_for_robot(s, img, painting.cpu().numpy(), allowed_colors, ignore_whites=True, output_dir='coolin_paint_instructions/'+str(i))\n",
    "    i += 1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}