Skip to content

Commit

Permalink
Add tracking mode options: Implemented fields to allow to choose the …
Browse files Browse the repository at this point in the history 
…desired tracking mode
  • Loading branch information
@jellydreams
jellydreams committed Dec 11, 2024
1 parent ba11f07 commit dfb7388
Show file tree
Hide file tree
Showing 3 changed files with 120 additions and 56 deletions.
13 changes: 9 additions & 4 deletions app.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

from info import VERSION
from plugin.ui import window_tracking_configuration, NIZIMA_LIVE, VTUBE_STUDIO
from plugin.mediapipe import draw_landmarks_on_image, get_bodyparts_values, init_mediapipe_options
from plugin.mediapipe import get_bodyparts_values, MediapipeTracking, LIVE_STREAM, IMAGE
from plugin.vtube_studio import connection_vts, create_parameters_vts, send_paramters_vts

from plugin.nizima import Nizima
Expand All @@ -30,7 +30,8 @@

async def main(settings):
# ----- MEDIAPIPE: LANDMARKER CONFIGURATION -----------
options = init_mediapipe_options(model_path)
mt = MediapipeTracking(mode=settings['tracking_mode'])
options = mt.init_mediapipe_options(model_path)
PoseLandmarker = mp.tasks.vision.PoseLandmarker

connection = False
Expand Down Expand Up @@ -91,15 +92,19 @@ async def main(settings):
timestamp += 1
# Detect pose landmarks from the current frame
input_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
RESULT = landmarker.detect(input_image)
if mt.mode == LIVE_STREAM:
landmarker.detect_async(input_image, timestamp)
RESULT = mt.result
else:
RESULT = landmarker.detect(input_image) # IMAGE

# Display Image for tracking window
image = render_image(input_image, settings['preview_enabled'])
if RESULT:
if RESULT.pose_world_landmarks:
# Get coordinates
parameters = RESULT
image = draw_landmarks_on_image(image, RESULT)
image = mt.draw_landmarks_on_image(image, RESULT)
else:
error_pose_estimation(image)

Expand Down
110 changes: 67 additions & 43 deletions plugin/mediapipe.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from mediapipe.tasks import python
from enum import Enum

LIVE_STREAM = 'LIVE_STREAM'
IMAGE = 'IMAGE'

class BodyParts(Enum):
"""
Expand Down Expand Up @@ -94,49 +96,71 @@ def get_part_from_name(i):
raise None


def init_mediapipe_options(model_path):
# Create a PoseLandmarker object
BaseOptions = mp.tasks.BaseOptions
PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
VisionRunningMode = mp.tasks.vision.RunningMode

options = PoseLandmarkerOptions(
base_options=BaseOptions(model_asset_path=model_path),
running_mode=VisionRunningMode.IMAGE)

return options


def draw_landmarks_on_image(img, detection_result):
"""
Draw landmarks on the input image.
:param rgb_image: input image
:param detection_result: result of landmark detection
:param preview: Whether to display the original image
:param annotated: Whether to annotate landmarks with their coordinates
:return: Image with landmarks
"""

# Fetch the image coordinates of the pose landmarks for drawing the pose on the image
pose_landmarks_list = detection_result.pose_landmarks

# - Loop through the detected poses to visualize
for idx in range(len(pose_landmarks_list)):
pose_landmarks = pose_landmarks_list[idx]

# -- Draw the pose landmarks
pose_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
pose_landmarks_proto.landmark.extend([
landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in pose_landmarks
])
solutions.drawing_utils.draw_landmarks(
img,
pose_landmarks_proto,
solutions.pose.POSE_CONNECTIONS,
solutions.drawing_styles.get_default_pose_landmarks_style())

return img
class MediapipeTracking:
def __init__(self, mode=LIVE_STREAM):
self.result = None
self.options = None
self.mode = mode

def init_mediapipe_options(self, model_path):
# Create a PoseLandmarker object
BaseOptions = mp.tasks.BaseOptions
PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions
VisionRunningMode = mp.tasks.vision.RunningMode

if self.mode == LIVE_STREAM:
PoseLandmarkerResult = mp.tasks.vision.PoseLandmarkerResult

def get_result(result: PoseLandmarkerResult, output_image: mp.Image, timestamp_ms: int):
print('pose landmarker result: {}'.format(result))
self.result = result

self.options = PoseLandmarkerOptions(
base_options=BaseOptions(model_asset_path=model_path),
running_mode=VisionRunningMode.LIVE_STREAM,
result_callback=get_result # LIVE_STREAM
)
else:
self.options = PoseLandmarkerOptions(
base_options=BaseOptions(model_asset_path=model_path),
running_mode=VisionRunningMode.IMAGE,
)
return self.options

def pose_detection(self, input_image, timestamp=None):
# Detect pose landmarks from the current frame
pass

def draw_landmarks_on_image(self, img, detection_result):
"""
Draw landmarks on the input image.
:param rgb_image: input image
:param detection_result: result of landmark detection
:param preview: Whether to display the original image
:param annotated: Whether to annotate landmarks with their coordinates
:return: Image with landmarks
"""

# Fetch the image coordinates of the pose landmarks for drawing the pose on the image
pose_landmarks_list = detection_result.pose_landmarks

# - Loop through the detected poses to visualize
for idx in range(len(pose_landmarks_list)):
pose_landmarks = pose_landmarks_list[idx]

# -- Draw the pose landmarks
pose_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
pose_landmarks_proto.landmark.extend([
landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in pose_landmarks
])
solutions.drawing_utils.draw_landmarks(
img,
pose_landmarks_proto,
solutions.pose.POSE_CONNECTIONS,
solutions.drawing_styles.get_default_pose_landmarks_style())

return img


def get_parameters_names():
Expand Down
53 changes: 44 additions & 9 deletions plugin/ui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,15 @@
import platform
import subprocess

from plugin.mediapipe import LIVE_STREAM, IMAGE

BACKGROUND_COLOR = "#333333"
TEXT_COLOR = "white"

BACKGROUND_VTUBE_STUDIO = '#fb58bb'
BACKGROUND_NIZIMA = '#692fc9'

WINDOW_SIZE = "340x440"
WINDOW_SIZE = "340x540"

NIZIMA_LIVE = 'NizimaLIVE'
VTUBE_STUDIO = 'VTube Studio'
Expand All @@ -38,12 +40,14 @@ def get_configuration():
api_port = 22022
else:
api_port = 8001
tracking_mode = selected_tracking_mode.get()

settings = {
'camera_url': camera_url,
'preview_enabled': show_camera_view_var.get(),
'port': api_port,
'software': software
'software': software,
'tracking_mode': tracking_mode
}

if available_cameras:
Expand Down Expand Up @@ -76,6 +80,15 @@ def toggle_inputs():
vtube_port_label.config(fg="gray")
vtube_name_label.config(fg="gray")

def toggle_inputs_tracking_mode():
"""Enable or disable the input fields based on the selected software."""
if selected_tracking_mode.get() == LIVE_STREAM:
mode_live_stream_label.config(fg="white")
mode_image_label.config(fg="gray")
elif selected_tracking_mode.get() == IMAGE:
mode_image_label.config(fg="white")
mode_live_stream_label.config(fg="gray")

root = tk.Tk()
root.title(f"VTS Fullbody Tracking {VERSION} - Settings")
root.geometry(WINDOW_SIZE)
Expand Down Expand Up @@ -138,17 +151,14 @@ def toggle_inputs():
vtube_studio_frame.pack(fill=tk.X, pady=(0, 5))

# --------------
# Use ttk to add styling to button
style = ttk.Style(root)
style.configure('TButton', bg='skyblue', fg='white')

# Chargement des icônes
# load vtube studio icon
image = os.path.abspath(os.path.join(os.path.dirname(__file__), ICON_VTUBE_STUDIO))
image = Image.open(image.replace("plugin\\", ""))
image_resized = image.resize((25, 25))
vtube_icon = ImageTk.PhotoImage(image_resized)
# Charger l'image avec Pillow

# load nizima icon
image = os.path.abspath(os.path.join(os.path.dirname(__file__), ICON_NIZIMA))
image = Image.open(image.replace("plugin\\", ""))
image_resized = image.resize((25, 25))
Expand All @@ -166,7 +176,6 @@ def toggle_inputs():
vtube_row.pack(anchor="w", pady=5)

vtube_radio = tk.Radiobutton(vtube_row, variable=selected_software, value=VTUBE_STUDIO, command=toggle_inputs, bg=BACKGROUND_COLOR )

vtube_radio.pack(side="left", padx=5)

vtube_icon_label = tk.Label(vtube_row, image=vtube_icon, bg=BACKGROUND_COLOR, width=18)
Expand Down Expand Up @@ -202,13 +211,39 @@ def toggle_inputs():
nizima_port_input.pack(side="left", padx=5)
nizima_port_input.insert(0, "22022") # Default port

# Initialiser les états des champs
# Init state fields
toggle_inputs()

# -- Tracking mode
tracking_setting_frame = create_section_header(root, "Tracking Settings")
tracking_setting_frame.pack(fill=tk.X, pady=(0, 5))

# Variable to track the selected software
selected_tracking_mode = tk.StringVar(value=LIVE_STREAM)

tracking_mode_frame = tk.Frame(root, padx=30, bg=BACKGROUND_COLOR)
tracking_mode_frame.pack(anchor="w", pady=5)

live_stream_row = tk.Frame(tracking_mode_frame, bg=BACKGROUND_COLOR)
live_stream_row.pack(anchor="w", pady=5)
mode_live_stream_radio = tk.Radiobutton(live_stream_row, variable=selected_tracking_mode, value=LIVE_STREAM, command=toggle_inputs_tracking_mode, bg=BACKGROUND_COLOR)
mode_live_stream_radio.pack(side="left", padx=5)
mode_live_stream_label = tk.Label(live_stream_row, text="Smooth and Slow", bg=BACKGROUND_COLOR, fg=TEXT_COLOR, anchor="w")
mode_live_stream_label.pack(side="left", padx=5)

image_row = tk.Frame(tracking_mode_frame, bg=BACKGROUND_COLOR)
image_row.pack(anchor="w", pady=5)
mode_image_radio = tk.Radiobutton(image_row, variable=selected_tracking_mode, value=IMAGE, command=toggle_inputs_tracking_mode, bg=BACKGROUND_COLOR)
mode_image_radio.pack(side="left", padx=5)
mode_image_label = tk.Label(image_row, text="Fast and Shaky", bg=BACKGROUND_COLOR, fg=TEXT_COLOR, anchor="w")
mode_image_label.pack(side="left", padx=5)

# -- Start Tracking Button
start_tracking_button = tk.Button(root, text="Start Tracking", command=root.quit, font=('Arial', 14, 'bold'), bg='#07121d', fg='white', activebackground='#3c9fbb', activeforeground='white', bd=0)
start_tracking_button.pack(pady=(5, 20), padx=20, fill=tk.X)

toggle_inputs_tracking_mode()

root.mainloop()

tracking_configuration = get_configuration()
Expand Down

0 comments on commit dfb7388

Please sign in to comment.