{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "24431280-e4f0-4c08-8e08-2d91814e48de", "metadata": { "tags": [] }, "outputs": [], "source": [ "import os\n", "import json\n", "import numpy as np\n", "from PIL import Image\n", "import matplotlib.pyplot as plt\n", "import matplotlib.patches as patches\n", "import cv2\n", "\n", "import natsort" ] }, { "cell_type": "code", "execution_count": 2, "id": "d799ad4d-9dea-4766-90d2-c8e1da51f724", "metadata": { "tags": [] }, "outputs": [], "source": [ "def get_class_color_mapping(annotations):\n", " classes = {category['id']: category['name'] for category in annotations['categories']}\n", "\n", " class_color_mapping = {}\n", " for class_id, class_name in classes.items():\n", " if class_id == 1:\n", " color = (0, 0, 1, 1) # Blue\n", " elif class_id == 2:\n", " color = (1, 0, 0, 1) # Red\n", " elif class_id == 30311:\n", " color = (0, 1, 0, 1) # Green\n", " elif class_id == 30319:\n", " color = (0, 0, 1, 1) # Blue\n", " else:\n", " color = plt.cm.tab20(class_id - 1)\n", " class_color_mapping[class_id] = color\n", "\n", " return class_color_mapping\n", "\n", "def save_annotated_image(image_path, label_path, output_path, class_color_mapping):\n", " image = Image.open(image_path)\n", " with open(label_path, 'r') as f:\n", " annotations = json.load(f)\n", "\n", " fig, ax = plt.subplots()\n", " ax.imshow(image)\n", "\n", " for annotation in annotations['annotations']:\n", " segmentation = annotation['segmentation'][0]\n", " x, y = zip(*[(segmentation[i], segmentation[i + 1]) for i in range(0, len(segmentation), 2)])\n", " class_id = annotation['category_id']\n", " color = class_color_mapping[class_id]\n", " polygon = patches.Polygon(np.column_stack((x, y)), edgecolor=color, facecolor=(*color[:3], 0.3))\n", " ax.add_patch(polygon)\n", "\n", " plt.axis('off')\n", " plt.savefig(output_path, bbox_inches='tight', pad_inches=0)\n", " plt.close()\n", "\n", "def create_video_from_images(images_path, labels_path, output_video):\n", " image_files = sorted([f for f in os.listdir(images_path) if f.endswith('.jpg')])\n", "\n", " if not os.path.exists('temp'):\n", " os.mkdir('temp')\n", "\n", " video_writer = None\n", " for idx, image_file in enumerate(image_files):\n", " label_file = os.path.splitext(image_file)[0] + '.json'\n", " # print(f\"Processing {image_file} with annotations from {label_file}\")\n", "\n", " image_path = os.path.join(images_path, image_file)\n", " label_path = os.path.join(labels_path, label_file)\n", "\n", " image = Image.open(image_path)\n", " with open(label_path, 'r') as f:\n", " annotations = json.load(f)\n", "\n", " class_color_mapping = get_class_color_mapping(annotations)\n", "\n", " output_path = os.path.join('temp', f\"annotated_{image_file[:-4]}.png\")\n", " save_annotated_image(image_path, label_path, output_path, class_color_mapping)\n", "\n", " if video_writer is None:\n", " height, width, _ = cv2.imread(output_path).shape\n", " video_writer = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*'mp4v'), 5, (width,height))\n", "\n", " annotated_image = cv2.imread(output_path)\n", " video_writer.write(annotated_image)\n", "\n", " video_writer.release()\n", "\n", " # Clean up the temporary annotated images\n", " for image_file in os.listdir('temp'):\n", " if image_file.endswith('.png'):\n", " os.remove(os.path.join('temp', image_file))\n", " os.rmdir('temp')" ] }, { "cell_type": "code", "execution_count": 3, "id": "fd72837d-69f7-4cc7-859e-a719dc4c1908", "metadata": { "tags": [] }, "outputs": [ { "ename": "OSError", "evalue": "[Errno 39] Directory not empty: 'temp'", "output_type": "error", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mOSError\u001b[0m Traceback (most recent call last)", "Cell \u001b[0;32mIn[3], line 5\u001b[0m\n\u001b[1;32m 2\u001b[0m labels_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m./dataset/label\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m 3\u001b[0m output_video \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124moutput_video.mp4\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m----> 5\u001b[0m \u001b[43mcreate_video_from_images\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimages_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlabels_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moutput_video\u001b[49m\u001b[43m)\u001b[49m\n", "Cell \u001b[0;32mIn[2], line 76\u001b[0m, in \u001b[0;36mcreate_video_from_images\u001b[0;34m(images_path, labels_path, output_video)\u001b[0m\n\u001b[1;32m 74\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m image_file\u001b[38;5;241m.\u001b[39mendswith(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m.png\u001b[39m\u001b[38;5;124m'\u001b[39m):\n\u001b[1;32m 75\u001b[0m os\u001b[38;5;241m.\u001b[39mremove(os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtemp\u001b[39m\u001b[38;5;124m'\u001b[39m, image_file))\n\u001b[0;32m---> 76\u001b[0m \u001b[43mos\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrmdir\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mtemp\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n", "\u001b[0;31mOSError\u001b[0m: [Errno 39] Directory not empty: 'temp'" ] } ], "source": [ "images_path = './dataset/image'\n", "labels_path = './dataset/label'\n", "output_video = 'output_video.mp4'\n", "\n", "create_video_from_images(images_path, labels_path, output_video)" ] }, { "cell_type": "code", "execution_count": null, "id": "1823a065-84a4-4e0d-ba8d-f83106e8231c", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "mmyolo", "language": "python", "name": "mmyolo" }, "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.8.16" } }, "nbformat": 4, "nbformat_minor": 5 }