get_cutouts_by_prompt.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# set up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install pycocotools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install torch torchvision torchaudio -f https://download.pytorch.org/whl/cu121/torch_stable.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response for image 1: {\"task_id\":\"1bd68d03-eb13-4245-94cb-fddde4a80f67\"}\n",
      "Response for image 2: {\"task_id\":\"712457fa-7b5c-4bd6-9c32-669fd884d3db\"}\n",
      "Response for image 3: {\"task_id\":\"b670a1e3-4678-4551-be5d-f31c8474b47b\"}\n",
      "Response for image 4: {\"task_id\":\"b59030bc-c09d-4cac-be64-8e65a38079a8\"}\n",
      "Response for image 5: {\"task_id\":\"3bf43508-7713-4c8d-91b6-c134938cd92a\"}\n",
      "Response for image 6: {\"task_id\":\"e0cb0caf-e842-445d-90fc-3c972133fd7a\"}\n",
      "Response for image 7: {\"task_id\":\"5903b4ad-c4ee-4630-8d09-48bfb8241206\"}\n",
      "Response for image 8: {\"task_id\":\"f00e1f5e-0eef-401c-8a56-bd60061d67bc\"}\n",
      "Response for image 9: {\"task_id\":\"db399403-3406-4ee0-a45f-5f14c8fd5bd0\"}\n",
      "Response for image 10: {\"task_id\":\"c318e628-92f9-4526-9f47-954a72d5c01c\"}\n",
      "Response for image 11: {\"task_id\":\"fe10ff75-b45f-4363-a8d5-87ec9e55515b\"}\n",
      "Response for image 12: {\"task_id\":\"5816d4cd-8163-4d60-9386-a7802bfb7bd5\"}\n",
      "Response for image 13: {\"task_id\":\"03e71391-bb07-4418-99af-111e99b8c25c\"}\n",
      "Response for image 14: {\"task_id\":\"dd7b8ab3-7bf6-4755-b299-cbe801da1fc3\"}\n",
      "Response for image 15: {\"task_id\":\"c3dae5da-ff9d-4edf-8761-d0194319f86b\"}\n",
      "Response for image 16: {\"task_id\":\"c4baa14b-d82a-46c1-bba0-a918a559252f\"}\n",
      "Response for image 17: {\"task_id\":\"66ea2e79-4280-48a2-971d-59322147fa02\"}\n",
      "Response for image 18: {\"task_id\":\"9c860b63-27e5-4ed3-afff-151b332ec9ce\"}\n",
      "Response for image 19: {\"task_id\":\"18a188fe-5b49-4228-a88a-329dafa8e898\"}\n",
      "Response for image 20: {\"task_id\":\"e556e9e4-48e9-41d0-9791-508be1538ebe\"}\n",
      "Response for image 21: {\"task_id\":\"e8038cf8-f291-444a-9c9d-6eb18094d526\"}\n",
      "Response for image 22: {\"task_id\":\"dcb6690a-93ba-4828-9bd4-2e51b385163e\"}\n",
      "Response for image 23: {\"task_id\":\"456154a3-ec7f-4d2e-86e7-6918fa2a7779\"}\n",
      "Response for image 24: {\"task_id\":\"d56feef0-7fcc-4039-abc1-c2349a7b43ad\"}\n",
      "Response for image 25: {\"task_id\":\"2898c40f-6f81-4e66-bb02-180241f640e7\"}\n",
      "Response for image 26: {\"task_id\":\"1187123d-4061-4043-9275-d6122430b83e\"}\n",
      "Response for image 27: {\"task_id\":\"e9938cf7-85bf-4be9-b811-5802d307a853\"}\n",
      "Response for image 28: {\"task_id\":\"4887ebf6-6f61-469e-867a-bc341db951b4\"}\n",
      "Response for image 29: {\"task_id\":\"0d57c9cf-7d43-4273-ba83-4f89b72be5c8\"}\n",
      "Response for image 30: {\"task_id\":\"e449a311-f10c-483f-9a24-a03d70ecce42\"}\n",
      "Response for image 31: {\"task_id\":\"e7d8e73c-5d4f-4379-aaab-56c047cee510\"}\n",
      "Response for image 32: {\"task_id\":\"cb9ec177-c2ef-4a65-9db3-56c89ad45566\"}\n",
      "Response for image 33: {\"task_id\":\"b8f3e99d-285e-4280-809f-200fe7bae740\"}\n",
      "Response for image 34: {\"task_id\":\"97e58b01-d66b-4803-8210-581b55d106b0\"}\n",
      "Response for image 35: {\"task_id\":\"700e929b-ada6-4624-b36b-af258e0b39c0\"}\n",
      "Response for image 36: {\"task_id\":\"7dc92f7b-7ef3-4dc5-8b10-56e8c471a57a\"}\n",
      "Response for image 37: {\"task_id\":\"bc89d421-becc-4142-8176-6f10b5333b90\"}\n",
      "Response for image 38: {\"task_id\":\"7c1e3052-7dd2-42d1-89e6-41210ec05ef2\"}\n",
      "Response for image 39: {\"task_id\":\"50350a7f-c976-4bd9-b761-5ff8bc254320\"}\n",
      "Response for image 40: {\"task_id\":\"2568d810-0ea8-4c51-9c85-332556b5307d\"}\n",
      "Response for image 41: {\"task_id\":\"7753fd92-2bfe-45cd-8994-f34644d82112\"}\n",
      "Response for image 42: {\"task_id\":\"4ba525aa-ee9c-4656-a94e-9794c98f9d86\"}\n",
      "Response for image 43: {\"task_id\":\"a77db0e7-6b69-49c4-916a-21454e2bcb9e\"}\n",
      "Response for image 44: {\"task_id\":\"66b54348-d7f4-44f8-8481-ba2fe7460fd5\"}\n",
      "Response for image 45: {\"task_id\":\"359ebb7f-fa12-46f5-8895-afc34b94a8aa\"}\n",
      "Response for image 46: {\"task_id\":\"209ea33d-acce-4f77-82b5-56651c3ba21f\"}\n",
      "Response for image 47: {\"task_id\":\"14321aeb-cc3c-446d-bfb2-63494006fa8f\"}\n",
      "Response for image 48: {\"task_id\":\"62b3d3a4-6fce-4b4d-9b10-428dce27aeee\"}\n",
      "Response for image 49: {\"task_id\":\"b3713991-141c-44a0-80a4-63970992ab4d\"}\n",
      "Response for image 50: {\"task_id\":\"944b3b9a-0da2-48d8-acc2-708e0901d975\"}\n",
      "Response for image 51: {\"task_id\":\"e348b9a4-6565-432f-bdb7-fd167bf12057\"}\n",
      "Response for image 52: {\"task_id\":\"6b586937-2ecc-4220-91e6-977bd3cfe77b\"}\n",
      "Response for image 53: {\"task_id\":\"de9ad433-5fff-47a2-b49d-c7b1e22a5540\"}\n",
      "Response for image 54: {\"task_id\":\"6c6c9120-e223-4aa7-a3ed-3d444d7c8140\"}\n",
      "Response for image 55: {\"task_id\":\"428d472c-fb00-4db1-bd7a-3f1a41edd229\"}\n",
      "Response for image 56: {\"task_id\":\"548ff2ba-6c23-4267-9377-665a16ba25cc\"}\n",
      "Response for image 57: {\"task_id\":\"12b1b12b-8803-4b6a-94a1-4ce0a8a617f2\"}\n",
      "Response for image 58: {\"task_id\":\"727c038a-7b75-4a55-b8f0-5a3160e08a56\"}\n",
      "Response for image 59: {\"task_id\":\"5efa68d1-fb4f-4827-961c-72df725d9e4c\"}\n",
      "Response for image 60: {\"task_id\":\"8a86c47f-45b5-4bf8-a1bf-43ef974499c9\"}\n",
      "Response for image 61: {\"task_id\":\"cc92a7c7-45e3-48a6-b408-b167096e7fbb\"}\n",
      "Response for image 62: {\"task_id\":\"b26884a2-9d51-4854-9924-8629c7eb02e0\"}\n",
      "Response for image 63: {\"task_id\":\"c609b651-0c4e-42f9-b97f-404ef1e7aa75\"}\n",
      "Response for image 64: {\"task_id\":\"b50e8c82-aed5-469e-8e1b-6e48fdc5be1d\"}\n",
      "Response for image 65: {\"task_id\":\"3e6e5c31-f7ab-41e2-a5bc-324cd727db04\"}\n",
      "Response for image 66: {\"task_id\":\"b27076da-a2e7-4874-a9ca-b46cdf1db7c4\"}\n",
      "Response for image 67: {\"task_id\":\"328416de-afa1-4c60-9209-d4b65d4d7957\"}\n",
      "Response for image 68: {\"task_id\":\"d53f3f19-ba79-4c80-84a0-53b8d4f04fb5\"}\n",
      "Response for image 69: {\"task_id\":\"4acbf3b9-d0f9-483f-907e-59acc6a87b63\"}\n",
      "Response for image 70: {\"task_id\":\"d85adfe6-1627-420b-bdba-8d2c3770c732\"}\n",
      "Response for image 71: {\"task_id\":\"c15ed633-18c3-4801-a757-5c90fc2e7785\"}\n",
      "Response for image 72: {\"task_id\":\"9ee7227d-ee30-4d6f-a3a5-9c98a54db9fc\"}\n",
      "Response for image 73: {\"task_id\":\"1efd432b-ac11-4e56-87f9-e00925074633\"}\n",
      "Response for image 74: {\"task_id\":\"9c6033b5-72c7-429c-82db-b9dc5dc543e1\"}\n",
      "Response for image 75: {\"task_id\":\"a902cc07-13e0-40ee-9da1-49a7d2eaa956\"}\n",
      "Response for image 76: {\"task_id\":\"f38d34a0-7e0d-402d-a24f-bf11f7f119fc\"}\n",
      "Response for image 77: {\"task_id\":\"8e789a7e-2d57-4e4a-b1a7-f9efec5d13ce\"}\n",
      "Response for image 78: {\"task_id\":\"cfb6fe24-0be2-4485-a632-63d059aaa8f7\"}\n",
      "Response for image 79: {\"task_id\":\"f0eb1334-9887-4aee-91ba-3711fea7f50c\"}\n",
      "Response for image 80: {\"task_id\":\"e641e690-d628-4f0b-bc66-9658d7dfa4ab\"}\n",
      "Response for image 81: {\"task_id\":\"eb6635f8-833b-4a39-af0e-71d0e1f41946\"}\n",
      "Response for image 82: {\"task_id\":\"049ac2aa-63e9-47cd-86b2-b962127b9cc6\"}\n",
      "Response for image 83: {\"task_id\":\"8719f81f-79ba-440e-921c-0c3de1140bac\"}\n",
      "Response for image 84: {\"task_id\":\"e55c8e45-556d-4e37-b481-93578e1871c6\"}\n",
      "Response for image 85: {\"task_id\":\"ea19e33b-038c-44f2-816f-53b8f649880d\"}\n",
      "Response for image 86: {\"task_id\":\"53c39fb3-94d3-469f-915d-5fc7e220a685\"}\n",
      "Response for image 87: {\"task_id\":\"a250b136-92c2-4fb7-be30-3bf9e2e80fa5\"}\n",
      "Response for image 88: {\"task_id\":\"67e88483-517b-449b-9503-ad36db1829f4\"}\n",
      "Response for image 89: {\"task_id\":\"41a1ef9e-2492-42d0-bfab-c45bacb4e167\"}\n",
      "Response for image 90: {\"task_id\":\"d4d5e0dd-7aa6-4ae9-bf5f-5bf65b7393be\"}\n",
      "Response for image 91: {\"task_id\":\"8050cd06-4abe-4ca2-b589-8c8ca08361e4\"}\n",
      "Response for image 92: {\"task_id\":\"643cca32-3471-4368-abff-e9f0fec9797e\"}\n",
      "Response for image 93: {\"task_id\":\"1434216e-d35a-4639-b307-76463ae299f3\"}\n",
      "Response for image 94: {\"task_id\":\"b9e9efba-d8b2-4116-98c4-a316cc7f1000\"}\n",
      "Response for image 95: {\"task_id\":\"6778ac34-817a-4c27-a9a0-4adada24cb22\"}\n",
      "Response for image 96: {\"task_id\":\"d9b3ad6d-1826-4fa6-8212-e8dd8b6500df\"}\n",
      "Response for image 97: {\"task_id\":\"fc004bcd-2fb3-4d35-a83d-1fd9195a1d04\"}\n",
      "Response for image 98: {\"task_id\":\"d7231859-4600-4842-8519-e0649a785bcb\"}\n",
      "Response for image 99: {\"task_id\":\"fc1c8042-fb23-4e56-8734-b9cbc4325fe6\"}\n",
      "Response for image 100: {\"task_id\":\"a5cd52c6-b0d0-48aa-98dc-9fa5678a1ae8\"}\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import requests\n",
    "import base64\n",
    "import json\n",
    "\n",
    "images_folder = \"./Images\"\n",
    "prompt = \"person\"\n",
    "\n",
    "for i, filename in enumerate(os.listdir(images_folder)[:100]):\n",
    "    # Construct the path to the image file\n",
    "    image_path = os.path.join(images_folder, filename)\n",
    "\n",
    "    # Read the image file and encode it as base64\n",
    "    with open(image_path, \"rb\") as image_file:\n",
    "        encoded_image = base64.b64encode(image_file.read()).decode(\"utf-8\")\n",
    "\n",
    "    # Make the API call\n",
    "    url = \"https://apps.beam.cloud/oregk\"\n",
    "    payload = {\"image\": encoded_image, \"prompt\": prompt}\n",
    "    headers = {\n",
    "        \"Accept\": \"*/*\",\n",
    "        \"Accept-Encoding\": \"gzip, deflate\",\n",
    "        \"Authorization\": \"Basic NzU0NTA2YmNjYzE4ZjNiMjNiMTc3OTgwNTY3ODcyNzQ6MTc4NDc2MTNiYmUyNDc3OTZlMGZlYTY0ZTJkYzMzMDY=\",\n",
    "        \"Connection\": \"keep-alive\",\n",
    "        \"Content-Type\": \"application/json\",\n",
    "    }\n",
    "    response = requests.request(\"POST\", url, headers=headers, data=json.dumps(payload))\n",
    "\n",
    "    # Print the response\n",
    "    print(f\"Response for image {i+1}: {response.text}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import sys\n",
    "import os\n",
    "from imageai.Detection import ObjectDetection\n",
    "import torch\n",
    "import json\n",
    "import cv2\n",
    "import numpy as np\n",
    "sys.path.append(\"..\")\n",
    "import pycocotools.mask as maskUtils\n",
    "import subprocess\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_mask(mask, ax, random_color=False):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "    else:\n",
    "        color = np.array([30/255, 144/255, 255/255, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "    \n",
    "def show_points(coords, labels, ax, marker_size=375):\n",
    "    pos_points = coords[labels==1]\n",
    "    neg_points = coords[labels==0]\n",
    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   \n",
    "    \n",
    "def show_box(box, ax):\n",
    "    x0, y0 = box[0], box[1]\n",
    "    w, h = box[2] - box[0], box[3] - box[1]\n",
    "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# Get a list of all the image files in the folder\n",
    "image_folder = '/home/noahr/repos/personal_projects/img_to_svg/Images'\n",
    "image_files = [os.path.join(image_folder, f) for f in os.listdir(image_folder) if f.endswith('.jpg')]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://download.pytorch.org/models/retinanet_resnet50_fpn_coco-eeacb38b.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import sys\n",
    "from segment_anything import sam_model_registry, SamPredictor\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "# sam_checkpoint = \"sam_vit_h_4b8939.pth\"\n",
    "# model_type = \"vit_h\"\n",
    "\n",
    "sam_checkpoint = \"sam_vit_b_01ec64.pth\"\n",
    "model_type = \"vit_b\"\n",
    "\n",
    "device = \"cuda\"\n",
    "\n",
    "sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "sam.to(device=device)\n",
    "\n",
    "predictor = SamPredictor(sam)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define paths for folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the paths to the images and masks folders (Change these to other folders in your machine if you prefer that)\n",
    "import os\n",
    "\n",
    "# For Local\n",
    "images_path = os.path.join(os.getcwd(), 'Images')\n",
    "masks_path = os.path.join(os.getcwd(), 'masks')\n",
    "cutouts_path = os.path.join(os.getcwd(), 'cutouts')\n",
    "svg_path = os.path.join(os.getcwd(), 'svg')\n",
    "video_images_path = os.path.join(os.getcwd(), 'video_images')\n",
    "\n",
    "\n",
    "# For Beam\n",
    "# Define the paths to the images and masks folders (Change these to other folders in your machine if you prefer that)\n",
    "# images_path = os.path.join('uploaded_images/Images')\n",
    "# masks_path = os.path.join('masks')\n",
    "# cutouts_path = os.path.join('cutouts')\n",
    "# svg_path = os.path.join('svg')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate coco masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U ultralytics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-cls.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ultralytics import YOLO\n",
    "import cv2\n",
    "\n",
    "model = YOLO('yolov8n.pt')\n",
    "results = model('Images/00100sPORTRAIT_00100_BURST20221127125031811_COVER.jpg')\n",
    "img = cv2.imread('Images/00100sPORTRAIT_00100_BURST20221127125031811_COVER.jpg')\n",
    "for result in results:\n",
    "    boxes = result.boxes.cpu().numpy()\n",
    "    for i, box in enumerate(boxes):\n",
    "        r = box.xyxy[0].astype(int)\n",
    "        print(r)\n",
    "        crop = img[r[1]:r[3], r[0]:r[2]]\n",
    "        cv2.imwrite(str(i) + \".jpg\", crop)\n",
    "        print(\"names: \",result.names[int(box.cls[0])])\n",
    "        print(str(i) + \".jpg\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mutli-threaded masking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "from ultralytics import YOLO\n",
    "from pycocotools import mask as maskUtils\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "from queue import Queue\n",
    "\n",
    "# Load the YOLOv5 model from torchub\n",
    "model = YOLO('yolov8n.pt')\n",
    "# Define the class you want to detect\n",
    "class_to_detect = 'person'\n",
    "\n",
    "# Define the path to the log file\n",
    "log_file = os.path.join(os.getcwd(), 'yolo_mask.log')\n",
    "\n",
    "# Define the number of workers to use\n",
    "num_workers = 4\n",
    "\n",
    "# Define a function to process a single image\n",
    "def process_image(filename):\n",
    "    # Check if a mask file already exists for this image\n",
    "    mask_number = 0\n",
    "    mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(mask_number)\n",
    "    while os.path.exists(os.path.join(masks_path, mask_filename)):\n",
    "        mask_number += 1\n",
    "        mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(mask_number)\n",
    "    if mask_number > 0:\n",
    "        return\n",
    "    \n",
    "    # Load the image\n",
    "    image = cv2.imread(os.path.join(images_path, filename))\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    # Redirect the YOLO model output to the log file\n",
    "    sys.stdout = open(log_file, 'w')\n",
    "    \n",
    "    # Run the model on the image\n",
    "    results = model(image)\n",
    "    predictor.set_image(image)\n",
    "    \n",
    "    for result in results:\n",
    "        boxes = result.boxes.cpu().numpy()\n",
    "        for i, box in enumerate(boxes):\n",
    "            if result.names[int(box.cls[0])] == class_to_detect:\n",
    "                masks, scores, _ = predictor.predict(\n",
    "                    point_coords=None,\n",
    "                    point_labels=None,\n",
    "                    box=box.xyxy[0].astype(int),\n",
    "                    multimask_output=True,\n",
    "                )\n",
    "                print(\"scores: \", scores)\n",
    "                for i, mask in enumerate(masks):\n",
    "                    # Convert the mask numpy array to a binary mask\n",
    "                    mask_binary = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.uint8)\n",
    "                    mask_binary[mask > 0] = 1\n",
    "                        \n",
    "                    # Convert the binary mask to a COCO RLE format\n",
    "                    mask_rle = maskUtils.encode(np.asfortranarray(mask_binary))\n",
    "                        \n",
    "                    # Extract the counts key from the mask RLE dictionary\n",
    "                    counts = mask_rle['counts']\n",
    "                        \n",
    "                    # Create a new dictionary with the required keys for COCO RLE format\n",
    "                    mask_coco_rle = {\n",
    "                        'size': [mask.shape[0], mask.shape[1]],\n",
    "                        'counts': counts.decode('utf-8'),\n",
    "                    }\n",
    "                        \n",
    "                    # Save the mask COCO RLE to a file\n",
    "                    mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(i)\n",
    "                    with open(os.path.join(masks_path, mask_filename), 'w') as f:\n",
    "                        json.dump(mask_coco_rle, f)\n",
    "\n",
    "# Get a list of all the image filenames\n",
    "image_filenames = os.listdir(images_path)\n",
    "print(\"FileName: \", image_filenames)\n",
    "# Create a thread-safe queue to hold the image filenames\n",
    "queue = Queue()\n",
    "for filename in image_filenames:\n",
    "    queue.put(filename)\n",
    "\n",
    "# Define a function to process images from the queue\n",
    "def process_queue():\n",
    "    while not queue.empty():\n",
    "        filename = queue.get()\n",
    "        process_image(filename)\n",
    "        queue.task_done()\n",
    "\n",
    "# Create a thread pool with the specified number of workers\n",
    "with ThreadPoolExecutor(max_workers=num_workers) as executor:\n",
    "    # Submit each image filename to the thread pool\n",
    "    futures = [executor.submit(process_queue) for _ in range(num_workers)]\n",
    "\n",
    "    # Wait for all the futures to complete\n",
    "    for future in as_completed(futures):\n",
    "        # Print any exceptions that occurred during processing\n",
    "        if future.exception() is not None:\n",
    "            print(future.exception())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Normal single-file masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mCannot execute code, session has been disposed. Please try restarting the Kernel."
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import os\n",
    "from ultralytics import YOLO\n",
    "\n",
    "# Load the YOLOv5 model from torchub\n",
    "model = YOLO('yolov8n.pt')\n",
    "# Define the class you want to detect\n",
    "class_to_detect = 'person'\n",
    "\n",
    "# Loop through all the files in the images folder\n",
    "for filename in os.listdir(images_path):\n",
    "    # Check if a mask file already exists for this image\n",
    "    mask_number = 0\n",
    "    mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(mask_number)\n",
    "    while os.path.exists(os.path.join(masks_path, mask_filename)):\n",
    "        mask_number += 1\n",
    "        mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(mask_number)\n",
    "    if mask_number > 0:\n",
    "        continue\n",
    "    \n",
    "    # Load the image\n",
    "    image = cv2.imread(os.path.join(images_path, filename))\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    # Run the model on the image\n",
    "    results = model(image)\n",
    "    predictor.set_image(image)\n",
    "    \n",
    "    for result in results:\n",
    "        boxes = result.boxes.cpu().numpy()\n",
    "        for i, box in enumerate(boxes):\n",
    "            # print(\"names: \",result.names[int(box.cls[0])])\n",
    "            if result.names[int(box.cls[0])] == class_to_detect:\n",
    "                masks, scores, _ = predictor.predict(\n",
    "                    point_coords=None,\n",
    "                    point_labels=None,\n",
    "                    box=box.xyxy[0].astype(int),\n",
    "                    multimask_output=True,\n",
    "                )\n",
    "                print(\"scores: \", scores)\n",
    "                for i, mask in enumerate(masks):\n",
    "                    # Convert the mask numpy array to a binary mask\n",
    "                    mask_binary = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.uint8)\n",
    "                    mask_binary[mask > 0] = 1\n",
    "                        \n",
    "                    # Convert the binary mask to a COCO RLE format\n",
    "                    mask_rle = maskUtils.encode(np.asfortranarray(mask_binary))\n",
    "                        \n",
    "                    # Extract the counts key from the mask RLE dictionary\n",
    "                    counts = mask_rle['counts']\n",
    "                        \n",
    "                    # Create a new dictionary with the required keys for COCO RLE format\n",
    "                    mask_coco_rle = {\n",
    "                        'size': [mask.shape[0], mask.shape[1]],\n",
    "                        'counts': counts.decode('utf-8'),\n",
    "                    }\n",
    "                        \n",
    "                    # Save the mask COCO RLE to a file\n",
    "                    mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(i)\n",
    "                    with open(os.path.join(masks_path, mask_filename), 'w') as f:\n",
    "                        json.dump(mask_coco_rle, f)\n",
    "                "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create cutouts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "from pycocotools import mask as maskUtils\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "from queue import Queue\n",
    "\n",
    "# Loop through all the files in the images folder\n",
    "for filename in os.listdir(images_path):\n",
    "    # Check if there are any mask files in the folder\n",
    "    if not os.listdir(masks_path):\n",
    "        exit()\n",
    "    \n",
    "    # Load the image\n",
    "    image = cv2.imread(os.path.join(images_path, filename))\n",
    "    \n",
    "    # Load the corresponding mask\n",
    "    mask_number = 0\n",
    "    cutout_number = 0\n",
    "    mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(mask_number)\n",
    "    while os.path.exists(os.path.join(masks_path, mask_filename)):\n",
    "        with open(os.path.join(masks_path, mask_filename), 'r') as f:\n",
    "            rle_dict = json.load(f)\n",
    "        size = rle_dict['size']\n",
    "        counts = rle_dict['counts']\n",
    "        mask_decoded = maskUtils.decode({'size': size, 'counts': counts})\n",
    "        mask_binary = np.zeros((size[0], size[1]), dtype=np.uint8)\n",
    "        mask_binary[mask_decoded > 0] = 1\n",
    "\n",
    "        # Resize the mask to match the shape of the image\n",
    "        mask_resized = cv2.resize(mask_decoded, (image.shape[1], image.shape[0]))\n",
    "\n",
    "        # Extract the cutout from the image using the mask\n",
    "        cutout = image * mask_resized[..., np.newaxis]\n",
    "\n",
    "        # Create an alpha channel for the cutout image\n",
    "        alpha = np.zeros(cutout.shape[:2], dtype=np.uint8)\n",
    "        alpha[mask_resized > 0] = 255\n",
    "        cutout = cv2.merge((cutout, alpha))\n",
    "\n",
    "        # Crop the cutout image to the bounding rectangle\n",
    "        x, y, w, h = cv2.boundingRect(mask_resized)\n",
    "        cutout = cutout[y:y+h, x:x+w]\n",
    "\n",
    "        # Save the cutout to the cutouts folder\n",
    "        cutout_filename = os.path.splitext(filename)[0] + '_cutout_{}.png'.format(cutout_number)\n",
    "        if os.path.exists(os.path.join(cutouts_path, cutout_filename)):\n",
    "            os.remove(os.path.join(cutouts_path, cutout_filename))\n",
    "        cv2.imwrite(os.path.join(cutouts_path, cutout_filename), cutout)\n",
    "        \n",
    "        # Delete the mask file\n",
    "        os.remove(os.path.join(masks_path, mask_filename))\n",
    "        \n",
    "        # Increment the mask and cutout numbers\n",
    "        mask_number += 1\n",
    "        cutout_number += 1\n",
    "        mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(mask_number)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-threaded cutouts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "from pycocotools import mask as maskUtils\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "from queue import Queue\n",
    "\n",
    "# Define the number of workers to use\n",
    "num_workers = 10\n",
    "\n",
    "# Define a function to process a single image\n",
    "def process_image(filename):\n",
    "    # Check if there are any mask files in the folder\n",
    "    if not os.listdir(masks_path):\n",
    "        return\n",
    "    \n",
    "    # Load the image\n",
    "    image = cv2.imread(os.path.join(images_path, filename))\n",
    "    \n",
    "    # Load the corresponding mask\n",
    "    mask_number = 0\n",
    "    cutout_number = 0\n",
    "    mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(mask_number)\n",
    "    while os.path.exists(os.path.join(masks_path, mask_filename)):\n",
    "        with open(os.path.join(masks_path, mask_filename), 'r') as f:\n",
    "            rle_dict = json.load(f)\n",
    "        size = rle_dict['size']\n",
    "        counts = rle_dict['counts']\n",
    "        mask_decoded = maskUtils.decode({'size': size, 'counts': counts})\n",
    "        mask_binary = np.zeros((size[0], size[1]), dtype=np.uint8)\n",
    "        mask_binary[mask_decoded > 0] = 1\n",
    "\n",
    "        # Resize the mask to match the shape of the image\n",
    "        mask_resized = cv2.resize(mask_decoded, (image.shape[1], image.shape[0]))\n",
    "\n",
    "        # Extract the cutout from the image using the mask\n",
    "        cutout = image * mask_resized[..., np.newaxis]\n",
    "\n",
    "        # Create an alpha channel for the cutout image\n",
    "        alpha = np.zeros(cutout.shape[:2], dtype=np.uint8)\n",
    "        alpha[mask_resized > 0] = 255\n",
    "        cutout = cv2.merge((cutout, alpha))\n",
    "\n",
    "        # Crop the cutout image to the bounding rectangle\n",
    "        x, y, w, h = cv2.boundingRect(mask_resized)\n",
    "        cutout = cutout[y:y+h, x:x+w]\n",
    "\n",
    "        # Save the cutout to the cutouts folder\n",
    "        cutout_filename = os.path.splitext(filename)[0] + '_cutout_{}.png'.format(cutout_number)\n",
    "        if os.path.exists(os.path.join(cutouts_path, cutout_filename)):\n",
    "            os.remove(os.path.join(cutouts_path, cutout_filename))\n",
    "        cv2.imwrite(os.path.join(cutouts_path, cutout_filename), cutout)\n",
    "        \n",
    "        # Delete the mask file\n",
    "        os.remove(os.path.join(masks_path, mask_filename))\n",
    "        \n",
    "        # Increment the mask and cutout numbers\n",
    "        mask_number += 1\n",
    "        cutout_number += 1\n",
    "        mask_filename = os.path.splitext(filename)[0] + '_mask_{}.json'.format(mask_number)\n",
    "\n",
    "# Get a list of all the image filenames\n",
    "image_filenames = os.listdir(images_path)\n",
    "\n",
    "# Create a thread-safe queue to hold the image filenames\n",
    "queue = Queue()\n",
    "for filename in image_filenames:\n",
    "    queue.put(filename)\n",
    "\n",
    "# Define a function to process images from the queue\n",
    "def process_queue():\n",
    "    while not queue.empty():\n",
    "        filename = queue.get()\n",
    "        process_image(filename)\n",
    "        queue.task_done()\n",
    "\n",
    "# Create a thread pool with the specified number of workers\n",
    "with ThreadPoolExecutor(max_workers=num_workers) as executor:\n",
    "    # Submit each image filename to the thread pool\n",
    "    futures = [executor.submit(process_queue) for _ in range(num_workers)]\n",
    "\n",
    "    # Wait for all the futures to complete\n",
    "    for future in as_completed(futures):\n",
    "        # Print any exceptions that occurred during processing\n",
    "        if future.exception() is not None:\n",
    "            print(future.exception())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make svg from png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "# Define the paths to the cutouts and output folders\n",
    "cutouts_path = '/home/noahr/repos/personal_projects/img_to_svg/cutouts'\n",
    "output_path = '/home/noahr/repos/personal_projects/img_to_svg/svg'\n",
    "\n",
    "# Loop through all the files in the cutouts folder\n",
    "for filename in os.listdir(cutouts_path):\n",
    "    if filename.endswith('.png'):\n",
    "        # Define the paths to the input and output files\n",
    "        png_path = os.path.join(cutouts_path, filename)\n",
    "        pbm_path = os.path.join(cutouts_path, os.path.splitext(filename)[0] + '.pbm')\n",
    "        svg_filename = os.path.splitext(filename)[0] + '.svg'\n",
    "        svg_path = os.path.join(output_path, svg_filename)\n",
    "\n",
    "        # Convert the PNG to PBM format\n",
    "        with Image.open(png_path) as im:\n",
    "            im.convert('1').save(pbm_path)\n",
    "\n",
    "        # Call potrace to convert the PBM to SVG\n",
    "        subprocess.call(['potrace', '-s', pbm_path, '-o', svg_path])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get frames from video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install setuptools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install ez_setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install moviepy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Create detector object\n",
    "detector = ObjectDetection()\n",
    "\n",
    "# Set model path\n",
    "model_path = \"/home/noahr/repos/personal_projects/img_to_svg/retinanet_resnet50_fpn_coco-eeacb38b.pth\"\n",
    "detector.setModelTypeAsRetinaNet()\n",
    "detector.setModelPath(model_path)\n",
    "detector.loadModel()\n",
    "\n",
    "# Loop through all files in input folder\n",
    "for filename in os.listdir(images_path):\n",
    "    # Check if file is an mp4 video\n",
    "    if filename.endswith(\".mp4\"):\n",
    "        # Set video path\n",
    "        video_path = os.path.join(images_path, filename)\n",
    "        \n",
    "        # Extract frames from video\n",
    "        video_frames_path = os.path.join(video_images_path, filename)\n",
    "        os.makedirs(video_frames_path, exist_ok=True)\n",
    "        cap = cv2.VideoCapture(video_path)\n",
    "        frame_count = 0\n",
    "        while True:\n",
    "            ret, frame = cap.read()\n",
    "            if not ret:\n",
    "                break\n",
    "            frame_path = os.path.join(video_frames_path, f\"frame_{frame_count}.jpg\")\n",
    "            cv2.imwrite(frame_path, frame)\n",
    "            frame_count += 1\n",
    "        \n",
    "        # Loop through frames and detect objects\n",
    "        for frame in os.listdir(video_frames_path):\n",
    "            frame_path = os.path.join(video_frames_path, frame)\n",
    "            detections = detector.detectObjectsFromImage(input_image=frame_path, output_image_path=os.path.join(video_images_path, \"frame.jpg\"), minimum_percentage_probability=80)\n",
    "            \n",
    "            # Loop through detections and save images with people\n",
    "            for detection in detections:\n",
    "                if detection[\"name\"] == \"person\":\n",
    "                    image_path = os.path.join(video_images_path, f\"{filename}_{detection['box_points'][0]}_{detection['box_points'][1]}.jpg\")\n",
    "                    frame = cv2.imread(frame_path)\n",
    "                    cv2.imwrite(image_path, frame)\n",
    "\n",
    "# Delete all images and folders inside video_images_path\n",
    "for file in os.listdir(video_images_path):\n",
    "    file_path = os.path.join(video_images_path, file)\n",
    "    if os.path.isfile(file_path):\n",
    "        os.remove(file_path)\n",
    "    elif os.path.isdir(file_path):\n",
    "        os.rmdir(file_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## For WSL2 Remove all zone.identifier tags after C&P"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loop through files in folder\n",
    "for filename in os.listdir(images_path):\n",
    "    # Check if file has the zone.identifier extension\n",
    "    if filename.endswith(\".zone.identifier\"):\n",
    "        # Set file path\n",
    "        file_path = os.path.join(images_path, filename)\n",
    "        \n",
    "        # Delete file\n",
    "        os.remove(file_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make svg from png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "# Loop through all the files in the cutouts folder\n",
    "for filename in os.listdir(cutouts_path):\n",
    "    if filename.endswith('.png'):\n",
    "        # Define the paths to the input and output files\n",
    "        png_path = os.path.join(cutouts_path, filename)\n",
    "        pbm_path = os.path.join(cutouts_path, os.path.splitext(filename)[0] + '.pbm')\n",
    "        svg_filename = os.path.splitext(filename)[0] + '.svg'\n",
    "        svg_path = os.path.join(svg_path, svg_filename)\n",
    "\n",
    "        # Convert the PNG to PBM format\n",
    "        with Image.open(png_path) as im:\n",
    "            im.convert('1').save(pbm_path)\n",
    "\n",
    "        # Call potrace to convert the PBM to SVG\n",
    "        subprocess.call(['potrace', '-s', pbm_path, '-o', svg_path])"
   ]
  }
 ],
 "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.10.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}