Skip to content

Commit

Permalink
Add UDF demo notebook
Browse files Browse the repository at this point in the history
  • Loading branch information
@dominikWin
dominikWin authored Aug 28, 2024
1 parent 89fea90 commit 713b7d1
Showing 1 changed file with 300 additions and 0 deletions.
300 changes: 300 additions & 0 deletions misc/UDFs-demo.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,300 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "be540ffe-d4f2-4bb8-bf49-b15866712f22",
"metadata": {},
"outputs": [],
"source": [
"import cv2\n",
"import numpy as np\n",
"\n",
"def add_bottom_overlay(image, percent=20, alpha=0.5):\n",
" height, width, _ = image.shape\n",
" overlay_height = int(height * (percent / 100))\n",
" overlay = image.copy()\n",
" cv2.rectangle(overlay, (0, height - overlay_height), (width, height), (0, 0, 0), -1)\n",
" image[height - overlay_height:height, 0:width] = cv2.addWeighted(\n",
" overlay[height - overlay_height:height, 0:width], alpha, \n",
" image[height - overlay_height:height, 0:width], 1 - alpha, 0)\n",
"\n",
"def draw_text(image, text, position, font=cv2.FONT_HERSHEY_SIMPLEX, font_scale=0.8, color=(255, 255, 255), thickness=2):\n",
" \"\"\"Draw text on an image.\"\"\"\n",
" cv2.putText(image, text, position, font, font_scale, color, thickness, cv2.LINE_AA)\n",
"\n",
"def draw_bar(image, position, width, height, level, label, bar_bg_color=(100, 100, 100), bar_fg_color_start=(255, 255, 255), bar_fg_color_end=(150, 150, 150)):\n",
" \"\"\"Draw a bar with a gradient and rounded corners on an image.\"\"\"\n",
" def draw_rounded_rect(img, pt1, pt2, color, radius):\n",
" x1, y1 = pt1\n",
" x2, y2 = pt2\n",
" cv2.rectangle(img, (x1 + radius, y1), (x2 - radius, y2), color, -1)\n",
" cv2.rectangle(img, (x1, y1 + radius), (x2, y2 - radius), color, -1)\n",
" cv2.ellipse(img, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, -1)\n",
" cv2.ellipse(img, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, -1)\n",
" cv2.ellipse(img, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, -1)\n",
" cv2.ellipse(img, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, -1)\n",
"\n",
" def draw_gradient_bar(img, pt1, pt2, start_color, end_color):\n",
" x1, y1 = pt1\n",
" x2, y2 = pt2\n",
" bar_width = x2 - x1\n",
" bar_height = y2 - y1\n",
"\n",
" gradient_bar = np.zeros((bar_height, bar_width, 3), dtype=np.uint8)\n",
" for i in range(bar_width):\n",
" alpha = i / bar_width\n",
" color = tuple(int(start_color[j] * (1 - alpha) + end_color[j] * alpha) for j in range(3))\n",
" cv2.line(gradient_bar, (i, 0), (i, bar_height), color, 1)\n",
"\n",
" rounded_bar = np.zeros((bar_height, bar_width, 3), dtype=np.uint8)\n",
" draw_rounded_rect(rounded_bar, (0, 0), (bar_width, bar_height), (255, 255, 255), bar_height // 2)\n",
" \n",
" mask = cv2.cvtColor(rounded_bar, cv2.COLOR_BGR2GRAY)\n",
" mask_inv = cv2.bitwise_not(mask)\n",
"\n",
" img_bg = cv2.bitwise_and(img[y1:y2, x1:x2], img[y1:y2, x1:x2], mask=mask_inv)\n",
" bar_fg = cv2.bitwise_and(gradient_bar, gradient_bar, mask=mask)\n",
"\n",
" img[y1:y2, x1:x2] = cv2.add(img_bg, bar_fg)\n",
"\n",
" x, y = position\n",
" bar_level = int(width * level)\n",
"\n",
" # Bar background and foreground\n",
" draw_rounded_rect(image, (x, y), (x + width, y + height), bar_bg_color, height // 2)\n",
" draw_gradient_bar(image, (x, y), (x + bar_level, y + height), bar_fg_color_start, bar_fg_color_end)\n",
" draw_text(image, label, (x + width + 10, y + 15))\n",
"\n",
"def draw_circ_bool(img, state, pos, rad=12, color=(255, 255, 255)):\n",
" thickness = cv2.FILLED if state else 1\n",
" cv2.circle(img, pos, rad, color, thickness, cv2.LINE_AA)\n",
"\n",
"def draw_battery(image, position, width, height, level, border_color=(255, 255, 255), fill_color=(207, 255, 207), background_color=(20, 20, 20)):\n",
" \"\"\"\n",
" Draw a battery icon with the specified fill level.\n",
" \n",
" Parameters:\n",
" - image: The image on which to draw the battery.\n",
" - position: A tuple (x, y) representing the top-left corner of the battery icon.\n",
" - width: The width of the battery icon.\n",
" - height: The height of the battery icon.\n",
" - level: The fill level of the battery (0 to 1).\n",
" - border_color: The color of the battery border.\n",
" - fill_color: The color of the battery fill.\n",
" - background_color: The color of the battery background.\n",
" \"\"\"\n",
" x, y = position\n",
" # Draw the main battery rectangle\n",
" cv2.rectangle(image, (x, y), (x + width, y + height), border_color, 2)\n",
" \n",
" # Draw the positive terminal\n",
" terminal_width = int(width * 0.08)\n",
" cv2.rectangle(image, (x + width, y + int(height * 0.3)), (x + width + terminal_width, y + int(height * 0.7)), border_color, -1)\n",
" \n",
" # Draw the battery background\n",
" cv2.rectangle(image, (x + 2, y + 2), (x + width - 2, y + height - 2), background_color, -1)\n",
" \n",
" # Draw the filled part of the battery\n",
" fill_width = int((width - 4) * level)\n",
" cv2.rectangle(image, (x + 2, y + 2), (x + 2 + fill_width, y + height - 2), fill_color, -1)\n",
"\n",
"import cv2\n",
"import numpy as np\n",
"\n",
"def draw_line_graph(image, data, position, graph_size, axis_color=(255, 255, 255), line_color=(0, 255, 255), thickness=2):\n",
" \"\"\"\n",
" Draw a line graph on an image.\n",
"\n",
" Parameters:\n",
" - image: The image on which to draw the graph.\n",
" - data: A list of normalized data points (values between -1 and 1).\n",
" - position: A tuple (x, y) representing the top-left corner of the graph.\n",
" - graph_size: A tuple (width, height) representing the size of the graph.\n",
" - axis_color: The color of the graph axes.\n",
" - line_color: The color of the graph line.\n",
" - thickness: The thickness of the graph line.\n",
" \"\"\"\n",
" x, y = position\n",
" width, height = graph_size\n",
"\n",
" # Draw horizontal axis\n",
" cv2.line(image, (x, y + height // 2), (x + width, y + height // 2), axis_color, thickness)\n",
" \n",
" # Draw graph line\n",
" num_points = len(data)\n",
" step = width // (num_points - 1)\n",
" \n",
" for i in range(1, num_points):\n",
" pt1 = (x + (i - 1) * step, y + height // 2 - int(data[i - 1] * (height // 2)))\n",
" pt2 = (x + i * step, y + height // 2 - int(data[i] * (height // 2)))\n",
" cv2.line(image, pt1, pt2, line_color, thickness)\n",
"\n",
" # Draw vertical axis\n",
" cv2.line(image, (x, y), (x, y + height), axis_color, thickness)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d4bc3cef-fc6f-42a5-bbe0-44975335580c",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import cv2\n",
"\n",
"def draw_axes(frame, angles, center=(320, 240), length=100):\n",
" \"\"\"\n",
" Draws 3D axes on the given frame.\n",
"\n",
" Parameters:\n",
" frame: The image frame on which to draw the axes.\n",
" angles: A tuple or list of three angles (roll, pitch, yaw) in degrees.\n",
" center: The center point for the axes (default is (320, 240) for a 640x480 frame).\n",
" length: The length of each axis line (default is 100).\n",
" \"\"\"\n",
" roll, pitch, yaw = np.deg2rad(angles)\n",
"\n",
" # Rotation matrices\n",
" Rx = np.array([\n",
" [1, 0, 0],\n",
" [0, np.cos(roll), -np.sin(roll)],\n",
" [0, np.sin(roll), np.cos(roll)]\n",
" ])\n",
"\n",
" Ry = np.array([\n",
" [np.cos(pitch), 0, np.sin(pitch)],\n",
" [0, 1, 0],\n",
" [-np.sin(pitch), 0, np.cos(pitch)]\n",
" ])\n",
"\n",
" Rz = np.array([\n",
" [np.cos(yaw), -np.sin(yaw), 0],\n",
" [np.sin(yaw), np.cos(yaw), 0],\n",
" [0, 0, 1]\n",
" ])\n",
"\n",
" # Combined rotation matrix\n",
" R = Rz @ Ry @ Rx\n",
"\n",
" # Define the axes in 3D space\n",
" axes = np.array([\n",
" [length, 0, 0], # X-axis (red)\n",
" [0, length, 0], # Y-axis (green)\n",
" [0, 0, length] # Z-axis (blue)\n",
" ])\n",
"\n",
" # Project the 3D axes to 2D\n",
" axes_2d = np.dot(axes, R.T).astype(int)\n",
"\n",
" # Define colors for the axes\n",
" colors = [(0, 0, 255), (0, 255, 0), (255, 0, 0)]\n",
"\n",
" for axis, color in zip(axes_2d, colors):\n",
" pt = (center[0] + axis[0], center[1] - axis[1])\n",
" cv2.line(frame, center, pt, color, 2, cv2.LINE_AA)\n",
"\n",
" return frame\n",
"\n",
"import pandas as pd\n",
"data = {\n",
" 'roll': np.linspace(0, 360, 100),\n",
" 'pitch': np.sin(np.linspace(0, 2*np.pi, 100)) * 45,\n",
" 'yaw': np.cos(np.linspace(0, 2*np.pi, 100)) * 45\n",
"}\n",
"df = pd.DataFrame(data)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ec101273-0265-4382-8c9b-b16ac67b0b06",
"metadata": {},
"outputs": [],
"source": [
"import vidformer\n",
"import cv2\n",
"\n",
"server = vidformer.YrdenServer(bin=\"../target/release/vidformer-cli\")\n",
"# server = vidformer.YrdenServer(domain='localhost', port=5000)\n",
"\n",
"scale = vidformer.Filter(\"Scale\")\n",
"\n",
"class MyFilter(vidformer.UDF): \n",
" def filter(self, frame: vidformer.UDFFrame, i: int):\n",
" f = frame.data().copy()\n",
" height, width = f.shape[:2]\n",
" add_bottom_overlay(f, percent=25)\n",
"\n",
" speed, altitude, lox, ch4 = 3808 + i // 3, 43 + i // 34, max(0.8 - i * .001, 0.1), max(1.0 - i * .003, 0.1)\n",
" draw_text(f, f'SPEED {speed} KM/H', (200, height-150))\n",
" draw_text(f, f'ALTITUDE {altitude} KM', (200, height-110))\n",
"\n",
" # Draw the LOX and CH4 bars\n",
" draw_bar(f, (200, height-40), 200, 20, lox, 'LOX')\n",
" draw_bar(f, (200, height-80), 200, 20, ch4, 'CH4')\n",
"\n",
" draw_text(f, f'T+{i/24:0.2f}', (575, height-75), font_scale=2.0, thickness=4)\n",
"\n",
" draw_text(f, f'ALPHA', (945, height-90))\n",
" draw_circ_bool(f, False, (920, height-100))\n",
"\n",
" draw_text(f, f'GAMMA', (945, height-50))\n",
" draw_circ_bool(f, i % 50 < 25, (920, height-60))\n",
"\n",
" # draw_battery(f, (50, 50), 100, 40, 0.75)\n",
"\n",
" import math\n",
" num_points = 50\n",
" data = [0.6 * math.sin(2 * math.pi * 2 * x / num_points + i / 25) for x in range(num_points)]\n",
" position = (1080, height-150) # (x, y)\n",
" size = (150, 100) # (width, height)\n",
" # Draw the line graph on the image\n",
" draw_line_graph(f, data, position, size)\n",
"\n",
" draw_axes(f, df.iloc[i % len(df)].values, center=(100, height-100), length=80)\n",
"\n",
" return vidformer.UDFFrame(f, frame.frame_type())\n",
" \n",
" def filter_type(self, frame, text): \n",
" return frame\n",
" \n",
"mf_udf = MyFilter(\"MyFilter\") \n",
"mf = mf_udf.into_filter()\n",
"\n",
"tos = vidformer.Source(server, \"tos_720p\", \"tos_720p.mp4\", 0)\n",
"\n",
"domain = tos.ts()\n",
"\n",
"def render(t, i):\n",
" return scale(mf(scale(tos[t], format=\"rgb24\", width=1280, height=720), i), format=\"yuv420p\", width=1280, height=720)\n",
"\n",
"spec = vidformer.Spec(domain, render, tos.fmt())\n",
"# spec.play(server)\n",
"spec.save(server, \"tos-hud.mp4\")"
]
}
],
"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.11.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

0 comments on commit 713b7d1

Please sign in to comment.