skcd_stabdiff.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/eyaler/skcd/blob/main/skcd_stabdiff.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "<img src=\"https://user-images.githubusercontent.com/4436747/186605404-385dd5fd-45fb-471a-8200-b19c13926928.png\" alt=\"skcd\" align=\"left\" height=\"320\">\n",
        "\n",
        "# skcd\n",
        "\n",
        "### [scikit-xkcd]\n",
        "\n",
        "#### Automagic stick-figure animation\n",
        "\n",
        "#### By [Olivia Seow](https://www.oliviaseow.com) ([@ollywogs](https://twitter.com/ollywogs)\\) and [Eyal Gruss](https://eyalgruss.com) ([@eyaler](https://twitter.com/eyaler)\\)\n",
        "\n",
        "Made at [Stochastic Labs](http://stochasticlabs.org)\n",
        "\n",
        "https://github.com/eyaler/skcd\n",
        "\n",
        "<br><br><br><br>\n",
        "Powered by:\n",
        "<table align=\"center\">\n",
        "<tr align=\"center\">\n",
        "<td>\n",
        "<a href=\"https://xkcd.com\"><figure><img src=\"https://user-images.githubusercontent.com/4436747/186625256-6f759c3a-d19d-4d88-be1f-cab4a67513cd.png\" alt=\"xkcd\" height=\"128\"><br><br><figcaption>xkcd</figcaption></figure></a>\n",
        "</td>\n",
        "<td>\n",
        "<a href=\"https://github.com/danforthcenter/plantcv\"><figure><img src=\"https://user-images.githubusercontent.com/4436747/186595109-5d5cde8d-4801-4ebe-bbe7-e7167ab59acf.png\" alt=\"PlantCV\" height=\"128\"><br><br><figcaption>PlantCV</figcaption></figure></a>\n",
        "</td>\n",
        "<td>\n",
        "<a href=\"https://github.com/CompVis/stable-diffusion\"><figure><img src=\"https://user-images.githubusercontent.com/4436747/186594458-7f22c4ee-3d43-41e6-b187-37efc8f685d3.png\" alt=\"Stable Diffusion\" height=\"128\"><br><br><figcaption>Stable Diffusion</figcaption></figure></a>\n",
        "</td>\n",
        "<td>\n",
        "<a href=\"https://github.com/huggingface/diffusers\"><figure><img src=\"https://user-images.githubusercontent.com/4436747/186587153-ad8b9602-d749-4b83-9307-3bdb76d0fd77.png\" alt=\"Hugging Face Diffusers\" height=\"128\"><br><figcaption>Hugging Face<br>🧨 Diffusers</figcaption></figure></a>\n",
        "</td>\n",
        "<td>\n",
        "<a href=\"https://imaginarysoundscape.net\"><figure><img src=\"https://user-images.githubusercontent.com/4436747/186593722-06f0d51c-aecc-408e-9433-8d274eaf3e8a.png\" alt=\"Imaginary Soundscape\" height=\"128\"><br><br><figcaption>Imaginary Soundscape</figcaption></figure></a>\n",
        "</td>\n",
        "</tr>\n",
        "</table>\n"
      ],
      "metadata": {
        "id": "vycOxG6uEFuh"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6cPY9Ou4sWs_",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Setup\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive', force_remount=True)\n",
        "\n",
        "!pip install plantcv\n",
        "#!wget https://raw.githubusercontent.com/danforthcenter/plantcv/3063cde35c7c1e379f977baa1048ec881ff64978/plantcv/plantcv/morphology/segment_combine.py -O /usr/local/lib/python3.7/dist-packages/plantcv/plantcv/morphology/segment_combine.py\n",
        "from plantcv import plantcv as pcv\n",
        "\n",
        "\n",
        "from google.colab import files\n",
        "import os\n",
        "from google.colab.patches import cv2_imshow\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "from six.moves.urllib.request import urlopen\n",
        "from six import BytesIO\n",
        "\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "from PIL import ImageColor\n",
        "from PIL import ImageDraw\n",
        "from PIL import ImageFont\n",
        "from PIL import ImageOps\n",
        "\n",
        "import cv2\n",
        "from skimage.morphology import skeletonize, thin, medial_axis, opening, dilation\n",
        "from skimage.segmentation import watershed\n",
        "from collections import Counter\n",
        "import time\n",
        "import io\n",
        "from IPython import display\n",
        "\n",
        "\n",
        "module_handle = \"https://tfhub.dev/google/faster_rcnn/openimages_v4/inception_resnet_v2/1\"\n",
        "detector = hub.load(module_handle).signatures['default']\n",
        "\n",
        "\n",
        "!pip install diffusers==0.2.4\n",
        "!pip install transformers scipy ftfy\n",
        "from huggingface_hub import notebook_login\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Setup Stable Diffusion\n",
        "#@markdown You also need to accept the model license before downloading or using the weights. In this post we'll use model version `v1-4`, so you'll need to  visit [its card](https://huggingface.co/CompVis/stable-diffusion-v1-4), read the license and tick the checkbox if you agree. \n",
        "\n",
        "#@markdown You have to be a registered user in 🤗 Hugging Face Hub, and you'll also need to use an access token for the code to work. For more information on access tokens, please refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).\n",
        "\n",
        "#@markdown Run the cell to enter your token and log in \n",
        "\n",
        "notebook_login()\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "FrdxsIk_RbXB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Data\n",
        "#@markdown folder on google drive or leave empty to upload manually\n",
        "\n",
        "image_folder = 'skcd/examples'  #@param {type: \"string\"}\n",
        "\n",
        "#@markdown folder on google drive or leave empty to save on colab \n",
        "\n",
        "cache_folder = 'stable_diffusion_cache'  #@param {type: \"string\"}\n",
        "\n",
        "if image_folder:\n",
        "  image_folder = os.path.join('/content/gdrive/MyDrive', image_folder)\n",
        "else:\n",
        "  image_folder = '/content/skcd'\n",
        "  os.makedirs(image_folder, exist_ok=True)\n",
        "  %cd image_folder\n",
        "  uploaded = files.upload()\n",
        "\n",
        "if cache_folder:\n",
        "  cache_folder = os.path.join('/content/gdrive/MyDrive', cache_folder)\n",
        "\n"
      ],
      "metadata": {
        "id": "LA_hqb4d4TdM",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from diffusers import StableDiffusionPipeline\n",
        "\n",
        "# make sure you're logged in with `huggingface-cli login\n",
        "kwargs = dict(revision=\"fp16\", torch_dtype=torch.float16, use_auth_token=True)\n",
        "if cache_folder:\n",
        "  kwargs['cache_dir'] = cache_folder\n",
        "pipe = StableDiffusionPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", **kwargs)\n",
        "StableDiffusionPipeline.safety_checker = lambda x,y: x, 0\n",
        "pipe = pipe.to(\"cuda\")"
      ],
      "metadata": {
        "id": "WBMz8MkiTEL5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D9IwDpOtpIHW",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Detect\n",
        "output_folder = 'skcd/output_sd' #@param {type: \"string\"}\n",
        "limit_to_these_images =  None #@param   # 7,11,22,18\n",
        "take_prompt_from_filename = True #@param {type: \"boolean\"}\n",
        "prompt_suffix = 'background painting by Don Hertzfeldt'  #@param {type: \"string\"}\n",
        "text2img_seed = 0 #@param {type: \"integer\"}\n",
        "brightness_factor = 2 #@param {type: \"number\"} \n",
        "closing_steps =   1#@param {type: \"integer\"}\n",
        "max_boxes_to_show = 10 #@param {type: \"integer\"}\n",
        "score_threshold = 0.002 #@param {type: \"number\"}\n",
        "iou_threshold = 0.2 #@param {type: \"number\"}\n",
        "min_area_frac = 0 #@param {type: \"number\"}\n",
        "max_area_frac = 0.55 #@param {type: \"number\"}\n",
        "skeleton_method = 'lee' #@param ['zhang', 'lee', 'thin', 'medial']\n",
        "max_expand_comp_area_ratio =  3.5#@param {type: \"number\"}\n",
        "max_expand_total_area_ratio =  3.5#@param {type: \"number\"}\n",
        "ignore_classes = 'book', 'envelope', 'poster', 'whiteboard' #@param\n",
        "show_classes = 'boy', 'chair', 'clothing', 'furniture', 'girl', 'human', 'human body', 'mammal', 'man', 'person', 'plant', 'tableware', 'toy', 'tree', 'tripod', 'woman' #@param\n",
        "dilation_steps =   3#@param {type: \"integer\"}\n",
        "segregated_merge_classes =  'boy', 'chair', 'clothing', 'furniture', 'girl', 'head', 'human', 'human body', 'human face', 'human head', 'mammal', 'man', 'person', 'plant', 'tableware', 'toy', 'tree', 'tripod', 'woman'#@param\n",
        "min_part_size_pixels =  0#@param {type: \"number\"}\n",
        "min_part_size_height_frac =  0.02#@param {type: \"number\"}\n",
        "min_parts = 6#@param {type: \"integer\"}\n",
        "abs_eps=0 #@param {type: \"number\"}\n",
        "rel_eps=0 #@param {type: \"number\"}\n",
        "anim_frames =  120 #@param {type: \"integer\"}\n",
        "frame_ms =  40#@param {type: \"integer\"}\n",
        "left_hand_amp =  0.9,0#@param \n",
        "right_hand_amp =  0,1#@param \n",
        "left_hand_speed =  0.5,0#@param\n",
        "right_hand_speed =  0,0.5#@param\n",
        "left_hand_phase =  -1,0#@param\n",
        "right_hand_phase =  0,-1#@param\n",
        "x_speed =  2,-1.5#@param\n",
        "y_speed =  0#@param \n",
        "z_speed = 0.01,-0.001#@param   \n",
        "vantage_x_frac = 0.75  #@param {type: \"number\"}\n",
        "vantage_y_frac = 0.5  #@param {type: \"number\"}\n",
        "line_thickness =   0#@param {type: \"integer\"}\n",
        "\n",
        "if ignore_classes:\n",
        "  if isinstance(ignore_classes, str):\n",
        "    ignore_classes = [ignore_classes]\n",
        "  ignore_classes = [cls.lower().replace('_', '').replace('-', '') for cls in ignore_classes]\n",
        "\n",
        "if show_classes:\n",
        "  if isinstance(show_classes, str):\n",
        "    show_classes = [show_classes]\n",
        "  show_classes = [cls.lower().replace('_', '').replace('-', '') for cls in show_classes]\n",
        "\n",
        "if segregated_merge_classes:\n",
        "  if isinstance(segregated_merge_classes, str):\n",
        "    segregated_merge_classes = [segregated_merge_classes]\n",
        "  segregated_merge_classes = [cls.lower().replace('_', '').replace('-', '') for cls in segregated_merge_classes]\n",
        "\n",
        "\n",
        "output_folder = os.path.join('/content/gdrive/MyDrive', output_folder)\n",
        "os.makedirs(output_folder, exist_ok=True)\n",
        "\n",
        "if isinstance(limit_to_these_images, int):\n",
        "  limit_to_these_images = [limit_to_these_images]\n",
        "\n",
        "\n",
        "def draw_bounding_box_on_image(image,\n",
        "                               y1,\n",
        "                               x1,\n",
        "                               y2,\n",
        "                               x2,\n",
        "                               color,\n",
        "                               font,\n",
        "                               thickness=3,\n",
        "                               display_str_list=[]):\n",
        "  \"\"\"Adds a bounding box to an image.\"\"\"\n",
        "  text = Image.new(\"RGBA\", image.size, (255, 255, 255, 0))\n",
        "  draw = ImageDraw.Draw(text)\n",
        "  im_width, im_height = image.size\n",
        "  (left, right, top, bottom) = (x1 * im_width, x2 * im_width,\n",
        "                                y1 * im_height, y2 * im_height)\n",
        "  draw.line([(left, top), (left, bottom), (right, bottom), (right, top),\n",
        "             (left, top)],\n",
        "            width=thickness,\n",
        "            fill=color)\n",
        "\n",
        "  # If the total height of the display strings added to the top of the bounding\n",
        "  # box exceeds the top of the image, stack the strings below the bounding box\n",
        "  # instead of above.\n",
        "  display_str_heights = [font.getsize(ds)[1] for ds in display_str_list]\n",
        "  display_str_widths = [font.getsize(ds)[0] for ds in display_str_list]\n",
        "  # Each display_str has a top and bottom margin of 0.05x.\n",
        "  total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)\n",
        "  total_display_str_width = (1 + 2 * 0.05) * sum(display_str_widths)\n",
        "\n",
        "  text_bottom = min(bottom + total_display_str_height, im_height)\n",
        "  text_left = max(min(left, im_width - total_display_str_width), 0)\n",
        "  # Reverse list and print from bottom to top.\n",
        "  for display_str in display_str_list[::-1]:\n",
        "    text_width, text_height = font.getsize(display_str)\n",
        "    margin = np.ceil(0.05 * text_height)\n",
        "    draw.rectangle([(text_left, text_bottom - text_height - 2 * margin),\n",
        "                    (text_left + text_width, text_bottom)],\n",
        "                   fill=color)\n",
        "    draw.text((text_left + margin, text_bottom - text_height - margin),\n",
        "              display_str,\n",
        "              fill=\"black\",\n",
        "              font=font)\n",
        "    text_bottom -= text_height - 2 * margin\n",
        "  return Image.alpha_composite(image, text)\n",
        "\n",
        "from plantcv.plantcv import params\n",
        "from plantcv.plantcv import dilate\n",
        "from plantcv.plantcv import outputs\n",
        "from plantcv.plantcv import find_objects\n",
        "from plantcv.plantcv._debug import _debug\n",
        "\n",
        "def find_branch_pts(skel_img, mask=None, label=\"default\"):\n",
        "    \"\"\"Find branch points in a skeletonized image.\n",
        "    The branching algorithm was inspired by Jean-Patrick Pommier: https://gist.github.com/jeanpat/5712699\n",
        "\n",
        "    Inputs:\n",
        "    skel_img    = Skeletonized image\n",
        "    mask        = (Optional) binary mask for debugging. If provided, debug image will be overlaid on the mask.\n",
        "    label        = optional label parameter, modifies the variable name of observations recorded\n",
        "\n",
        "    Returns:\n",
        "    branch_pts_img = Image with just branch points, rest 0\n",
        "\n",
        "    :param skel_img: numpy.ndarray\n",
        "    :param mask: np.ndarray\n",
        "    :param label: str\n",
        "    :return branch_pts_img: numpy.ndarray\n",
        "    \"\"\"\n",
        "    # In a kernel: 1 values line up with 255s, -1s line up with 0s, and 0s correspond to don't care\n",
        "    # T like branch points\n",
        "    t1 = np.array([[-1, 1, -1],\n",
        "                   [1, 1, 1],\n",
        "                   [-1, 0, -1]])\n",
        "    t2 = np.array([[1, -1, 1],\n",
        "                   [-1, 1, -1],\n",
        "                   [1, -1, 0]])\n",
        "    t3 = np.rot90(t1)\n",
        "    t4 = np.rot90(t2)\n",
        "    t5 = np.rot90(t3)\n",
        "    t6 = np.rot90(t4)\n",
        "    t7 = np.rot90(t5)\n",
        "    t8 = np.rot90(t6)\n",
        "\n",
        "    # Y like branch points\n",
        "    y1 = np.array([[1, -1, 1],\n",
        "                   [0, 1, 0],\n",
        "                   [0, 1, 0]])\n",
        "    y2 = np.array([[-1, 1, -1],\n",
        "                   [1, 1, 0],\n",
        "                   [-1, 0, 1]])\n",
        "    y3 = np.rot90(y1)\n",
        "    y4 = np.rot90(y2)\n",
        "    y5 = np.rot90(y3)\n",
        "    y6 = np.rot90(y4)\n",
        "    y7 = np.rot90(y5)\n",
        "    y8 = np.rot90(y6)\n",
        "    \n",
        "   \n",
        "    kernels = [t1, t2, t3, t4, t5, t6, t7, t8, y1, y2, y3, y4, y5, y6, y7, y8]\n",
        "\n",
        "    branch_pts_img = np.zeros(skel_img.shape[:2], dtype=int)\n",
        "\n",
        "    # Store branch points\n",
        "    for kernel in kernels:\n",
        "        branch_pts_img = np.logical_or(cv2.morphologyEx(skel_img, op=cv2.MORPH_HITMISS, kernel=kernel,\n",
        "                                                        borderType=cv2.BORDER_CONSTANT, borderValue=0), branch_pts_img)\n",
        "\n",
        "    # Switch type to uint8 rather than bool\n",
        "    branch_pts_img = branch_pts_img.astype(np.uint8) * 255\n",
        "\n",
        "    # Store debug\n",
        "    debug = params.debug\n",
        "    params.debug = None\n",
        "\n",
        "    # Make debugging image\n",
        "    if mask is None:\n",
        "        dilated_skel = dilate(skel_img, params.line_thickness, 1)\n",
        "        branch_plot = cv2.cvtColor(dilated_skel, cv2.COLOR_GRAY2RGB)\n",
        "    else:\n",
        "        # Make debugging image on mask\n",
        "        mask_copy = mask.copy()\n",
        "        branch_plot = cv2.cvtColor(mask_copy, cv2.COLOR_GRAY2RGB)\n",
        "        skel_obj, skel_hier = find_objects(skel_img, skel_img)\n",
        "        cv2.drawContours(branch_plot, skel_obj, -1, (150, 150, 150), params.line_thickness, lineType=8,\n",
        "                         hierarchy=skel_hier)\n",
        "\n",
        "    branch_objects, _ = find_objects(branch_pts_img, branch_pts_img)\n",
        "\n",
        "    # Initialize list of tip data points\n",
        "    branch_list = []\n",
        "    branch_labels = []\n",
        "    for i, branch in enumerate(branch_objects):\n",
        "        x, y = branch.ravel()[:2]\n",
        "        coord = (int(x), int(y))\n",
        "        branch_list.append(coord)\n",
        "        branch_labels.append(i)\n",
        "        cv2.circle(branch_plot, (x, y), params.line_thickness, (255, 0, 255), -1)\n",
        "\n",
        "    outputs.add_observation(sample=label, variable='branch_pts',\n",
        "                            trait='list of branch-point coordinates identified from a skeleton',\n",
        "                            method='plantcv.plantcv.morphology.find_branch_pts', scale='pixels', datatype=list,\n",
        "                            value=branch_list, label=branch_labels)\n",
        "\n",
        "    # Reset debug mode\n",
        "    params.debug = debug\n",
        "\n",
        "    _debug(visual=branch_plot, filename=os.path.join(params.debug_outdir, f\"{params.device}_branch_pts.png\"))\n",
        "\n",
        "    return branch_pts_img\n",
        "\n",
        "\n",
        "def draw_boxes(image, boxes, classes, scores, merge_lists):\n",
        "  \"\"\"Overlay labeled boxes on an image with formatted scores and label names.\"\"\"\n",
        "  colors = [c for c in ImageColor.colormap.values() if isinstance(c, str) and 150 < np.mean(ImageColor.getrgb(c)) < 200]\n",
        "\n",
        "  try:\n",
        "    font = ImageFont.truetype(\"/usr/share/fonts/truetype/liberation/LiberationSansNarrow-Regular.ttf\",\n",
        "                              25)\n",
        "  except IOError:\n",
        "    print(\"Font not found, using default font.\")\n",
        "    font = ImageFont.load_default()\n",
        "\n",
        "  scores = np.array(scores)\n",
        "  boxes, classes, scores = zip(*sorted([x for x in zip(boxes, classes, scores)], key=lambda x: (show_classes and x[1] not in show_classes, area(x[0]), -x[2])))\n",
        "\n",
        "  bg = image.copy()\n",
        "  bg_colors = Counter(tuple(c) for c in bg.reshape(-1,3))\n",
        "  bg_color = bg_colors.most_common()[0][0]\n",
        "  print(bg_color)\n",
        "  results = []\n",
        "  for _ in range(closing_steps):\n",
        "    image = opening(image)\n",
        "  image = np.array(Image.fromarray(image).convert(\"RGBA\"))\n",
        "\n",
        "  im_height, im_width = image.shape[:2]\n",
        "  filter_classes = []\n",
        "  filter_scores = []\n",
        "  filter_merge_lists = []\n",
        "  masks = []\n",
        "  skeletons = []\n",
        "  objects = []\n",
        "  ind_skel_parts = []\n",
        "  tight_boxes = []\n",
        "  final_boxes = []\n",
        "  new_image = image.copy()\n",
        "  all_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n",
        "  for (y1, x1, y2, x2), cls, score, mlst in zip(boxes, classes, scores, merge_lists):\n",
        "    if score < score_threshold or (show_classes and cls not in show_classes):\n",
        "        continue\n",
        "    (left, right, top, bottom) = (int(x1 * im_width), int(x2 * im_width),\n",
        "                            int(y1 * im_height), int(y2 * im_height))\n",
        "    grab = image[top:bottom, left:right]\n",
        "    gray = cv2.cvtColor(grab, cv2.COLOR_RGB2GRAY)\n",
        "    th, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)\n",
        "    #binary = all_binary[top:bottom, left:right]\n",
        "    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(binary)\n",
        "    areas = [[stats[label, cv2.CC_STAT_AREA], label] for label in range(1, num_labels)]\n",
        "    a, max_label = sorted(areas)[-1]\n",
        "    mask = (labels == max_label).astype(np.uint8) * 255\n",
        "    x, y, w, h = cv2.boundingRect(mask)\n",
        "    box = (top+y) / im_height, (left+x) / im_width, (top+y+h) / im_height, (left+x+w) / im_width\n",
        "    if (max_expand_comp_area_ratio > 1 or not max_expand_comp_area_ratio) and (max_expand_total_area_ratio > 1 or not max_expand_total_area_ratio):\n",
        "      all_binary = (all_gray < th).astype(np.uint8) * 255     \n",
        "      num_all_labels, all_labels, all_stats, all_centroids = cv2.connectedComponentsWithStats(all_binary)\n",
        "      lst = all_labels[top:bottom, left:right] * (mask > 0)\n",
        "      lst = [l for l in lst.flatten() if l]\n",
        "      big_label = max(set(lst), key=lst.count)\n",
        "      big_mask = (all_labels == big_label).astype(np.uint8) * 255\n",
        "      x, y, w, h = cv2.boundingRect(big_mask)\n",
        "      final_box = box\n",
        "      big_box = y / im_height, x / im_width, (y+h) / im_height, (x+w) / im_width\n",
        "      if (all_stats[big_label, cv2.CC_STAT_AREA] <= a * max_expand_comp_area_ratio or not max_expand_comp_area_ratio) and (area(big_box) <= area(box) * max_expand_total_area_ratio or not max_expand_total_area_ratio):\n",
        "        final_box = big_box\n",
        "        mask = big_mask\n",
        "        top = left = 0\n",
        "        bottom = im_height\n",
        "        right = im_width\n",
        "    \n",
        "    have = False\n",
        "    for fb in final_boxes:\n",
        "      iou = bb_intersection_over_union(fb, final_box)\n",
        "      if iou >= iou_threshold:\n",
        "        have = True\n",
        "        break\n",
        "    if have:\n",
        "      continue\n",
        "    \n",
        "    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
        "    if skeleton_method == 'zhang':\n",
        "      sk = skeletonize(mask)\n",
        "    elif skeleton_method == 'lee':\n",
        "      sk = skeletonize(mask, method='lee')\n",
        "    elif skeleton_method == 'thin':\n",
        "      sk = thin(mask)\n",
        "    elif skeleton_method == 'medial':\n",
        "      sk = medial_axis(mask)\n",
        "    contour_mask = mask.copy()\n",
        "    cv2.drawContours(mask, contours, -1, color=255, thickness=cv2.FILLED)\n",
        "    overlay = new_image[top:bottom, left:right].copy()\n",
        "    overlay[mask > 0] = 255, 0, 0, 32\n",
        "    new_image[top:bottom, left:right] = np.array(Image.alpha_composite(Image.fromarray(new_image[top:bottom, left:right]), Image.fromarray(overlay)))\n",
        "    new_image[top:bottom, left:right][sk > 0] = 0, 255, 255, 255\n",
        "    dilated_mask = mask\n",
        "    for _ in range(dilation_steps):\n",
        "      dilated_mask = dilation(dilated_mask)\n",
        "    pcv.params.line_thickness = 2\n",
        "    pcv.params.color_scale = 'gist_ncar'\n",
        "    pcv.params.color_sequence = 'random'\n",
        "    #cycle_img = pcv.morphology.check_cycles(sk)\n",
        "    #num_cycles = pcv.outputs.observations['default']['num_cycles']['value']\n",
        "    bp = find_branch_pts(sk)\n",
        "    bp = pcv.dilate(bp, 3, 1)\n",
        "    segments = pcv.image_subtract(sk, bp)\n",
        "    segment_parts, _ = pcv.find_objects(segments, segments)\n",
        "    rand_color = pcv.color_palette(num=len(segment_parts), saved=False)\n",
        "    segmented_img = sk.copy()\n",
        "    segmented_img = cv2.cvtColor(segmented_img, cv2.COLOR_GRAY2RGB)\n",
        "    for q, cnt in enumerate(segment_parts):\n",
        "        cv2.drawContours(segmented_img, segment_parts, q, rand_color[q], thickness=1)        \n",
        "    \n",
        "    h, w = mask.shape\n",
        "    markers = np.zeros((h, w))\n",
        "    labels = np.arange(len(segment_parts)) + 1\n",
        "    for m, l in enumerate(labels):\n",
        "        cv2.drawContours(markers, segment_parts, m, color=int(l), thickness=1)\n",
        "    filled_mask = watershed(contour_mask == 0, markers=markers, mask=contour_mask != 0, compactness=0)\n",
        "    non_black = np.any(segmented_img != (0,0,0), axis=-1)\n",
        "    new_image[top:bottom, left:right, :3][non_black] = segmented_img[non_black]\n",
        "    skel_parts = []\n",
        "    small_skel_parts = []\n",
        "    other_parts = []\n",
        "    skel_holes = []\n",
        "    small_skel_holes = []\n",
        "    other_holes = []\n",
        "    \n",
        "    jj = 0\n",
        "    for j, part in enumerate(segment_parts, start=1):\n",
        "      this_mask = (filled_mask == j).astype(np.uint8)\n",
        "      contours, hierarchy = cv2.findContours(this_mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE) \n",
        "      arc = cv2.arcLength(part, closed=True) / 2\n",
        "      if arc < min_part_size_pixels or arc < (bottom-top) * min_part_size_height_frac:\n",
        "        continue\n",
        "      else:\n",
        "        jj += 1\n",
        "      if jj == 4:\n",
        "        segment_list = range(j + 1, len(segment_parts))\n",
        "        #if segment_list:\n",
        "        #  _, segment_parts = pcv.morphology.segment_combine(segment_list, segment_parts, mask)\n",
        "        break\n",
        "    \n",
        "    jj = 0\n",
        "    for j, part in enumerate(segment_parts, start=1):\n",
        "      this_mask = (filled_mask == j).astype(np.uint8)\n",
        "      contours, hierarchy = cv2.findContours(this_mask, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE) \n",
        "      if hierarchy is None:\n",
        "        continue\n",
        "      arc = cv2.arcLength(part, closed=True) / 2\n",
        "      if arc < min_part_size_pixels or arc < (bottom-top) * min_part_size_height_frac:\n",
        "        if line_thickness:\n",
        "          small_skel_parts.append((part, -1))\n",
        "        else:\n",
        "          for c,h in zip(contours, hierarchy[0]):\n",
        "            if h[3] == -1:\n",
        "              small_skel_parts.append((c, 0))\n",
        "            elif h[3] == 0:\n",
        "              small_skel_holes.append((c, 1))\n",
        "      elif jj > 4:\n",
        "        for c,h in zip(contours, hierarchy[0]):\n",
        "          if h[3] == -1:\n",
        "            other_parts.append((c, 0))\n",
        "          elif h[3] == 0:\n",
        "            other_holes.append((c, 1))\n",
        "      else:\n",
        "        if line_thickness:\n",
        "          approx = cv2.approxPolyDP(part, max(abs_eps, rel_eps * arc), closed=True)\n",
        "          skel_parts.append((approx, -1))\n",
        "        else:\n",
        "          for c,h in zip(contours, hierarchy[0]):\n",
        "            if h[3] == -1:\n",
        "              skel_parts.append((c, 0))\n",
        "            elif h[3] == 0:\n",
        "              skel_holes.append((c, 1))            \n",
        "        jj += 1\n",
        "    if len(skel_parts) + len(other_parts) >= min_parts:\n",
        "      tight_boxes.append(box)\n",
        "      final_boxes.append(final_box)\n",
        "      filter_classes.append(cls)\n",
        "      filter_scores.append(score)\n",
        "      filter_merge_lists.append(mlst)\n",
        "      masks.append(mask)\n",
        "      skeletons.append(sk)\n",
        "      bg[top:bottom, left:right][dilated_mask > 0] = bg_color\n",
        "      parts = []\n",
        "      for part, flag in small_skel_parts + small_skel_holes + other_parts + other_holes + skel_parts + skel_holes:\n",
        "        shifted = np.array([[[left + x[0][0], top + x[0][1]]] for x in part])\n",
        "        parts.append((shifted, flag))\n",
        "      objects.append(parts)\n",
        "      ind_skel_parts.append(len(small_skel_parts) + len(small_skel_holes) + len(other_parts) + len(other_holes))\n",
        "\n",
        "  image_pil = Image.fromarray(np.array(new_image))\n",
        "  classes, scores, merge_lists = filter_classes, filter_scores, filter_merge_lists\n",
        "  ordered = sorted(zip(tight_boxes, final_boxes, classes, scores, merge_lists, masks, skeletons, objects, ind_skel_parts), key=lambda x: (x[1][1] + x[1][3], x[1][0] + x[1][2]))\n",
        "  if ordered:\n",
        "    tight_boxes, final_boxes, classes, scores, merge_lists, masks, skeletons, objects, ind_skel_parts = zip(*ordered)\n",
        "  else:\n",
        "    objects = []\n",
        "\n",
        "  for i, (tight_box, final_box, cls, score, merge_list) in enumerate(zip(tight_boxes, final_boxes, classes, scores, merge_lists)):\n",
        "    if score >= score_threshold:       \n",
        "      display_str = \"{}. {} ({}%)\".format(i + 1, cls, round(100 * score, 1))\n",
        "      line_color = ImageColor.getcolor(colors[hash(cls) % len(colors)], 'RGB') + (192,)\n",
        "      if i < max_boxes_to_show:\n",
        "        if tight_box != final_box:\n",
        "          image_pil = draw_bounding_box_on_image(image_pil, *tight_box, line_color, font, thickness=1)\n",
        "          line_color = ImageColor.getcolor(colors[hash(cls) % len(colors)], 'RGB') + (192,)\n",
        "        image_pil = draw_bounding_box_on_image(image_pil, *final_box, line_color, font, display_str_list=[display_str])\n",
        "      y1, x1, y2, x2 = final_box\n",
        "      results.append(display_str +  f' {Counter(merge_list+[cls]).most_common()} {tight_box} {final_box} center_x={(x1+x2)/2*100:.0f}% center_y={(y1+y2)/2*100:.0f}% height={(y2-y1)*100:.0f}% area={area(final_box)*100:.1f}%')\n",
        "\n",
        "  return np.array(image_pil), results, tight_boxes, final_boxes, masks, skeletons, objects, ind_skel_parts, bg\n",
        "\n",
        "\n",
        "def load_image(path):\n",
        "  image = tf.io.read_file(path)\n",
        "  image = tf.image.decode_image(image, channels=3)\n",
        "  return image\n",
        "\n",
        "\n",
        "def area(box):\n",
        "  return (box[2]-box[0]) * (box[3]-box[1])\n",
        "\n",
        "\n",
        "def bb_intersection_over_union(boxA, boxB):\n",
        "    # determine the (x, y)-coordinates of the intersection rectangle\n",
        "    yA = max(boxA[0], boxB[0])\n",
        "    xA = max(boxA[1], boxB[1])\n",
        "    yB = min(boxA[2], boxB[2])\n",
        "    xB = min(boxA[3], boxB[3])\n",
        "\n",
        "    # compute the area of intersection rectangle\n",
        "    interArea = max((xB - xA, 0)) * max((yB - yA), 0)\n",
        "    if interArea == 0:\n",
        "        return 0\n",
        "    # compute the area of both the prediction and ground-truth\n",
        "    # rectangles\n",
        "    boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))\n",
        "    boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))\n",
        "\n",
        "    # compute the intersection over union by taking the intersection\n",
        "    # area and dividing it by the sum of prediction + ground-truth\n",
        "    # areas - the interesection area\n",
        "    iou = interArea / float(boxAArea + boxBArea - interArea)\n",
        "\n",
        "    # return the intersection over union value\n",
        "    return iou\n",
        "\n",
        "import copy\n",
        "def run_detector(detector, path):\n",
        "  try:\n",
        "    image = load_image(path)\n",
        "  except Exception:\n",
        "    return\n",
        "  converted_image = tf.image.convert_image_dtype(image, tf.float32)[tf.newaxis, ...]\n",
        "  start_time = time.time()\n",
        "  result = detector(converted_image)\n",
        "  end_time = time.time()\n",
        "\n",
        "  result = {key:value.numpy() for key,value in result.items()}\n",
        "\n",
        "  print(\"Found %d objects.\" % len(result[\"detection_scores\"]))\n",
        "  #print(\"Inference time: \", end_time-start_time)\n",
        "\n",
        "  boxes, classes, scores = result[\"detection_boxes\"], result[\"detection_class_entities\"], result[\"detection_scores\"]\n",
        "  bcs = list((b,c.decode().lower(),s) for b,c,s in zip(boxes, classes, scores) if (not ignore_classes or c.decode().lower() not in ignore_classes) and min_area_frac <= area(b) <= max_area_frac and s >= score_threshold)\n",
        "  if bcs:\n",
        "    boxes, classes, scores = zip(*bcs)\n",
        "\n",
        "  #indices = tf.image.non_max_suppression(\n",
        "  #  boxes,\n",
        "  #  scores,\n",
        "  #  max_boxes,\n",
        "  #  float(iou_threshold),\n",
        "  #  float(score_threshold),\n",
        "  #  name=None\n",
        "  #)\n",
        "  #boxes, classes, scores = zip(*[x for i, x in enumerate(zip(boxes, classes, scores)) if i in indices])\n",
        "\n",
        "  bcs = list(zip(boxes, classes, scores, [[] for _ in range(len(boxes))]))\n",
        "  while True:\n",
        "    best_iou = 0\n",
        "    for i in range(len(bcs) - 1):\n",
        "      for j in range(i + 1, len(bcs)):\n",
        "        iou = bb_intersection_over_union(bcs[i][0], bcs[j][0])\n",
        "        box = np.array([min(bcs[i][0][0], bcs[j][0][0]), min(bcs[i][0][1], bcs[j][0][1]), max(bcs[i][0][2], bcs[j][0][2]), max(bcs[i][0][3], bcs[j][0][3])])\n",
        "        if iou > best_iou and area(box) <= max_area_frac and (not segregated_merge_classes or (bcs[i][1] in segregated_merge_classes) == (bcs[j][1] in segregated_merge_classes)):\n",
        "          best_iou = iou\n",
        "          best_i = i\n",
        "          best_j = j\n",
        "    if best_iou >= iou_threshold:\n",
        "      bcs2 = bcs.pop(best_j)\n",
        "      bcs1 = bcs.pop(best_i)\n",
        "      box = np.array([min(bcs1[0][0], bcs2[0][0]), min(bcs1[0][1], bcs2[0][1]), max(bcs1[0][2], bcs2[0][2]), max(bcs1[0][3], bcs2[0][3])])\n",
        "      a1 = area(bcs1[0])\n",
        "      a2 = area(bcs2[0])\n",
        "      if a2 > a1 or a2 == a1 and bcs2[2] > bcs1[2]:\n",
        "        cls, score = bcs2[1:3]\n",
        "        other = bcs1[1]\n",
        "      else:\n",
        "        cls, score = bcs1[1:3]\n",
        "        other = bcs2[1]\n",
        "      bcs.append((box, cls, score, bcs1[3] + bcs2[3] + [other]))\n",
        "    else:\n",
        "      break\n",
        "  boxes, classes, scores, merge_lists = zip(*bcs)\n",
        "\n",
        "  image_with_boxes, results, tight_boxes, final_boxes, masks, skeletons, objects, ind_skel_parts, bg = draw_boxes(image.numpy(), boxes, classes, scores, merge_lists)\n",
        "  global left_hand_amp, right_hand_amp, left_hand_speed,right_hand_speed, left_hand_phase,right_hand_phase, x_speed, y_speed, z_speed\n",
        "  if isinstance(left_hand_amp, (int,float)):\n",
        "    left_hand_amp = [left_hand_amp]*len(objects)\n",
        "  if isinstance(right_hand_amp, (int,float)):\n",
        "    right_hand_amp = [right_hand_amp]*len(objects)\n",
        "  if isinstance(left_hand_speed, (int,float)):\n",
        "    left_hand_speed = [left_hand_speed]*len(objects)\n",
        "  if isinstance(right_hand_speed, (int,float)):\n",
        "    right_hand_speed = [right_hand_speed]*len(objects)\n",
        "  if isinstance(left_hand_phase, (int,float)):\n",
        "    left_hand_phase = [left_hand_phase]*len(objects)\n",
        "  if isinstance(right_hand_phase, (int,float)):\n",
        "    right_hand_phase = [right_hand_phase]*len(objects)\n",
        "  if isinstance(x_speed, (int,float)):\n",
        "    x_speed = [x_speed]*len(objects)\n",
        "  if isinstance(y_speed, (int,float)):\n",
        "    y_speed = [y_speed]*len(objects)\n",
        "  if isinstance(z_speed, (int,float)):\n",
        "    z_speed = [z_speed]*len(objects)\n",
        "\n",
        "  if len(left_hand_amp) < len(objects):\n",
        "    left_hand_amp += type(left_hand_amp)([0] * (len(objects)-len(left_hand_amp)))\n",
        "  if len(right_hand_amp) < len(objects):\n",
        "      right_hand_amp += type(right_hand_amp)([0] * (len(objects)-len(right_hand_amp)))\n",
        "  if len(left_hand_speed) < len(objects):\n",
        "      left_hand_speed += type(left_hand_speed)([0] * (len(objects)-len(left_hand_speed)))\n",
        "  if len(right_hand_speed) < len(objects):\n",
        "      right_hand_speed += type(right_hand_speed)([0] * (len(objects)-len(right_hand_speed)))\n",
        "  if len(left_hand_phase) < len(objects):\n",
        "      left_hand_phase += type(left_hand_phase)([0] * (len(objects)-len(left_hand_phase)))\n",
        "  if len(right_hand_phase) < len(objects):\n",
        "      right_hand_phase += type(right_hand_phase)([0] * (len(objects)-len(right_hand_phase)))\n",
        "  if len(x_speed) < len(objects):\n",
        "      x_speed += type(x_speed)([0] * (len(objects)-len(x_speed)))\n",
        "  if len(y_speed) < len(objects):\n",
        "      y_speed += type(y_speed)([0] * (len(objects)-len(y_speed)))\n",
        "  if len(z_speed) < len(objects):\n",
        "      z_speed += type(z_speed)([0] * (len(objects)-len(z_speed)))\n",
        "\n",
        "  cv_image = np.array(image_with_boxes)\n",
        "  cv_image[..., :3] = cv_image[..., 2::-1]\n",
        "  cv2_imshow(cv_image)\n",
        "  cv_bg = bg.copy()\n",
        "  cv_bg[..., :3] = cv_bg[..., 2::-1]\n",
        "  cv2_imshow(cv_bg)\n",
        "  putback = cv_bg.copy()\n",
        "  for o, parts in enumerate(objects):\n",
        "    parts, flags = zip(*parts)\n",
        "    for i, flag in enumerate(flags):\n",
        "      cv2.drawContours(putback, parts, i, color=(255,255,255) if flag == 1 else 0, thickness=line_thickness if flag == -1 else cv2.FILLED, lineType=cv2.LINE_AA)\n",
        "  cv2_imshow(putback)\n",
        "  frames = []\n",
        "  obj_th = {}\n",
        "  obj_r = {}\n",
        "  \n",
        "  from torch import autocast\n",
        "  from PIL import ImageEnhance\n",
        "\n",
        "  prompt = []\n",
        "  if take_prompt_from_filename:\n",
        "    prompt.append(os.path.splitext(os.path.basename(path))[0].replace('_', ' '))\n",
        "  if prompt_suffix:\n",
        "    prompt.append(prompt_suffix)\n",
        "  prompt = ', '.join(prompt)\n",
        "  print(prompt)\n",
        "  generator = None\n",
        "  div_factor = 64\n",
        "  if bg.shape[1] < bg.shape[0]:\n",
        "    width = 512\n",
        "    height = div_factor * round(width * bg.shape[1] / bg.shape[0] / div_factor)\n",
        "  else:\n",
        "    height = 512\n",
        "    width = div_factor * round(height * bg.shape[0] / bg.shape[1] / div_factor)\n",
        "  if text2img_seed:\n",
        "    generator = torch.Generator(\"cuda\").manual_seed(text2img_seed)\n",
        "  with autocast(\"cuda\"):\n",
        "    bg_img = pipe(prompt, generator=generator, height=height, width=width)[\"sample\"][0]\n",
        "  bg_img = bg_img.resize((bg.shape[1],bg.shape[0]))\n",
        "  enhancer = ImageEnhance.Brightness(bg_img)\n",
        "  bg_img = enhancer.enhance(brightness_factor) \n",
        "  bg_path = os.path.join(output_folder, os.path.splitext(os.path.basename(path))[0]+'_bg.jpg')\n",
        "  bg_img.save(bg_path)\n",
        "  bg_img = np.array(bg_img)\n",
        "  bg = np.minimum(bg_img, bg)\n",
        "    \n",
        "  for o, (parts, ind) in enumerate(zip(objects, ind_skel_parts)):\n",
        "    parts, flags = zip(*parts)\n",
        "    parts = list(parts)\n",
        "    avg_th = {0:0, 1:0}\n",
        "    max_r = {0:0, 1:0}\n",
        "    leg_diff = np.mean(parts[ind], axis=0)[0][0] > np.mean(parts[ind + 1], axis=0)[0][0]\n",
        "    for j in range(2):\n",
        "      part = copy.deepcopy(parts[ind + j : ind + j + 1])\n",
        "      if leg_diff and j==0 or not leg_diff and j==1:\n",
        "        ref = min(part[0], key=lambda x: (x[0][0], x[0][1]))\n",
        "      else:\n",
        "        ref = max(part[0], key=lambda x: (x[0][0], -x[0][1]))\n",
        "      part[0] = part[0] - ref\n",
        "      for k in range(len(part[0])):\n",
        "        r = np.sqrt(part[0][k][0][0]*part[0][k][0][0]+part[0][k][0][1]*part[0][k][0][1])\n",
        "        th = np.arctan2(part[0][k][0][0], part[0][k][0][1])\n",
        "        avg_th[j] += abs(th)\n",
        "        max_r[j] = max(max_r[j], r)\n",
        "      avg_th[j] = avg_th[j] / len(part[0])\n",
        "    obj_th[o] = (avg_th[0] + avg_th[1]) / len(parts[ind : ind + 2])\n",
        "    obj_r[o] = (max_r[0] + max_r[1])/len(parts[ind : ind + 2])\n",
        "    print(o, obj_r[o], obj_th[o]*180/np.pi)\n",
        "  midx = cv_image.shape[1] * vantage_x_frac\n",
        "  midy = cv_image.shape[0] * vantage_y_frac\n",
        "  for i in range(anim_frames):\n",
        "    frames.append(bg.copy())\n",
        "    z_order = sorted([(j,z_speed[j]) for j in range(len(objects))], key=lambda x:-x[1]) \n",
        "    for o, _ in z_order:\n",
        "      parts, flags = zip(*objects[o])\n",
        "      parts = list(parts)\n",
        "      ind = ind_skel_parts[o]\n",
        "      z_index = 1 + z_speed[o] * i\n",
        "      if z_index <= 0:\n",
        "        continue\n",
        "      for j in range(len(parts)):\n",
        "        part = copy.deepcopy(parts[j:j+1])\n",
        "        for k in range(len(part[0])):\n",
        "          part[0][k][0][0] = np.round((part[0][k][0][0] + i*x_speed[o] - midx) / z_index + midx)\n",
        "          part[0][k][0][1] = np.round((part[0][k][0][1] + i*y_speed[o] - midy) / z_index + midy)\n",
        "        \n",
        "        jj = j - ind\n",
        "\n",
        "        if len(parts) > ind + 4:\n",
        "          hand_diff = np.mean(parts[ind + 3], axis=0)[0][0] > np.mean(parts[ind + 4], axis=0)[0][0]\n",
        "          if jj in (3, 4):\n",
        "            if hand_diff and jj==3 or not hand_diff and jj==4:\n",
        "              rev = 1\n",
        "              ref = min(part[0], key=lambda x: (x[0][0], x[0][1]))\n",
        "              hand_amp = right_hand_amp\n",
        "              hand_speed = right_hand_speed\n",
        "              hand_phase = right_hand_phase\n",
        "            else:\n",
        "              rev = -1\n",
        "              ref = max(part[0], key=lambda x: (x[0][0], -x[0][1]))\n",
        "              hand_amp = left_hand_amp\n",
        "              hand_speed = left_hand_speed\n",
        "              hand_phase = left_hand_phase\n",
        "            part[0] = part[0] - ref\n",
        "            for k in range(len(part[0])):\n",
        "              r = np.sqrt(part[0][k][0][0]*part[0][k][0][0]+part[0][k][0][1]*part[0][k][0][1])\n",
        "              th = np.arctan2(part[0][k][0][0], part[0][k][0][1])\n",
        "              angle = th + hand_amp[o] * rev * (np.sin(hand_speed[o] * i + hand_phase[o]*np.pi/2) - np.sin(np.pi*(1 - hand_phase[o]/2)))\n",
        "              part[0][k][0][0] = np.round(r * np.sin(angle))\n",
        "              part[0][k][0][1] = np.round(r * np.cos(angle))\n",
        "            part[0] = part[0] + ref\n",
        "\n",
        "        fix_factor = np.pi  # to be removed\n",
        "        if len(parts) > ind + 1:\n",
        "          leg_diff = np.mean(parts[ind], axis=0)[0][0] > np.mean(parts[ind + 1], axis=0)[0][0]\n",
        "          if jj in (0, 1):\n",
        "            if leg_diff and jj==0 or not leg_diff and jj==1:\n",
        "              ref = min(part[0], key=lambda x: (x[0][0], x[0][1]))\n",
        "            else:\n",
        "              ref = max(part[0], key=lambda x: (x[0][0], -x[0][1]))\n",
        "            part[0] = part[0] - ref\n",
        "            for k in range(len(part[0])):\n",
        "              r = np.sqrt(part[0][k][0][0]*part[0][k][0][0]+part[0][k][0][1]*part[0][k][0][1])\n",
        "              th = np.arctan2(part[0][k][0][0], part[0][k][0][1])\n",
        "              angle = th * np.cos(fix_factor*x_speed[o]/(obj_r[o]*obj_th[o]) * i) if obj_r[o]*obj_th[o] != 0 else th\n",
        "              part[0][k][0][0] = np.round(r * np.sin(angle))\n",
        "              part[0][k][0][1] = np.round(r * np.cos(angle))\n",
        "            part[0] = part[0] + ref\n",
        "            \n",
        "        cv2.drawContours(frames[-1], part, 0, color=(255,255,255) if flags[j] == 1 else 0, thickness=line_thickness if flags[j] == -1 else cv2.FILLED, lineType=cv2.LINE_AA)\n",
        "\n",
        "    frames[-1] = Image.fromarray(frames[-1])\n",
        "  \n",
        "  out = os.path.join(output_folder, os.path.splitext(os.path.basename(path))[0]+'.gif')\n",
        "  frames[0].save(out, format='GIF', append_images=frames[1:], save_all=True, optimize=True, duration=frame_ms, loop=0)\n",
        "  \n",
        "  import requests\n",
        "  from requests.adapters import HTTPAdapter, Retry\n",
        "\n",
        "  s = requests.Session()\n",
        "  retries = Retry(total=5, backoff_factor=1)\n",
        "  s.mount('https://', HTTPAdapter(max_retries=retries))\n",
        "\n",
        "  with open(bg_path, 'rb') as f:\n",
        "    files = {'file': f.read()}\n",
        "  url = 'https://imaginarysoundscape.net:8000/process'\n",
        "  s = requests.post(url, files=files)\n",
        "  name = s.json()['sound_id']\n",
        "  name = f'https://storage.imaginarysoundscape.net/sounds/{name}.mp3'\n",
        "  audio = os.path.join(output_folder, os.path.splitext(os.path.basename(path))[0]+'.mp3')\n",
        "  !wget \"$name\" -O \"$audio\"\n",
        "  display.display(display.Audio(audio))\n",
        "  \n",
        "  display.display(display.Image(out))\n",
        "\n",
        "  for r in results:\n",
        "    print(r)\n",
        "  print()\n",
        "\n",
        "for i, fn in enumerate(sorted(os.listdir(image_folder))):\n",
        "  if limit_to_these_images and i not in limit_to_these_images:\n",
        "    continue\n",
        "  print(i, fn)\n",
        "  run_detector(detector, os.path.join(image_folder, fn))\n",
        "\n",
        "print('finished')\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "skcd_stabdiff",
      "private_outputs": true,
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}