666 lines
25 KiB
Python
666 lines
25 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from fastvideo.distributed.parallel_state import get_local_torch_device
|
|
from fastvideo.utils import logger
|
|
|
|
try:
|
|
import cv2
|
|
except ImportError:
|
|
cv2 = None
|
|
|
|
CAM_VALUE = 0.1
|
|
CAMERA_MAP = {"i": [CAM_VALUE, 0], "k": [-CAM_VALUE, 0], "j": [0, -CAM_VALUE], "l": [0, CAM_VALUE], "u": [0, 0]}
|
|
|
|
KEYBOARD_MAP_4 = { # base_distilled_model (universal): W/S/A/D
|
|
"w": [1, 0, 0, 0],
|
|
"s": [0, 1, 0, 0],
|
|
"a": [0, 0, 1, 0],
|
|
"d": [0, 0, 0, 1],
|
|
"q": [0, 0, 0, 0]
|
|
}
|
|
KEYBOARD_MAP_2 = { # gta_distilled_model: W/S only (steering via mouse)
|
|
"w": [1, 0],
|
|
"s": [0, 1],
|
|
"q": [0, 0]
|
|
}
|
|
KEYBOARD_MAP_7 = { # templerun_distilled_model: still/w/s/left/right/a/d
|
|
"q": [1, 0, 0, 0, 0, 0, 0], # still
|
|
"w": [0, 1, 0, 0, 0, 0, 0], # forward
|
|
"s": [0, 0, 1, 0, 0, 0, 0], # back
|
|
"j": [0, 0, 0, 1, 0, 0, 0], # left (swipe)
|
|
"l": [0, 0, 0, 0, 1, 0, 0], # right (swipe)
|
|
"a": [0, 0, 0, 0, 0, 1, 0], # a
|
|
"d": [0, 0, 0, 0, 0, 0, 1], # d
|
|
}
|
|
KEYBOARD_MAP = KEYBOARD_MAP_4 # Default for backward compatibility
|
|
SOLARIS_MOVEMENT_KEY_INDICES = {
|
|
"W": 11,
|
|
"S": 12,
|
|
"A": 13,
|
|
"D": 14,
|
|
}
|
|
|
|
|
|
def retain_kv_with_sink(
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
target_len: int,
|
|
sink_size: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if k.shape[1] <= target_len or sink_size <= 0:
|
|
return k[:, -target_len:], v[:, -target_len:]
|
|
|
|
sink_len = min(int(sink_size), target_len, k.shape[1])
|
|
tail_len = target_len - sink_len
|
|
sink_k = k[:, :sink_len]
|
|
sink_v = v[:, :sink_len]
|
|
if tail_len <= 0:
|
|
return sink_k, sink_v
|
|
|
|
tail_k = k[:, sink_len:][:, -tail_len:]
|
|
tail_v = v[:, sink_len:][:, -tail_len:]
|
|
return torch.cat([sink_k, tail_k], dim=1), torch.cat([sink_v, tail_v], dim=1)
|
|
|
|
|
|
def _require_cv2() -> bool:
|
|
if cv2 is None:
|
|
logger.warning("OpenCV is not available; skipping MatrixGame2 validation overlay.")
|
|
return False
|
|
return True
|
|
|
|
|
|
def _matrixgame2_overlay_keys_from_keyboard(keyboard_frame: np.ndarray | torch.Tensor, ) -> dict[str, bool]:
|
|
vec = np.asarray(keyboard_frame, dtype=np.float32).reshape(-1)
|
|
dim = vec.shape[0]
|
|
|
|
if dim >= 23:
|
|
return {key: bool(vec[idx] > 0.5) for key, idx in SOLARIS_MOVEMENT_KEY_INDICES.items()}
|
|
if dim == 7:
|
|
return {
|
|
"W": bool(vec[1] > 0.5),
|
|
"S": bool(vec[2] > 0.5),
|
|
"A": bool(vec[5] > 0.5),
|
|
"D": bool(vec[6] > 0.5),
|
|
}
|
|
if dim >= 6:
|
|
return {
|
|
"W": bool(vec[0] > 0.5),
|
|
"S": bool(vec[1] > 0.5),
|
|
"A": bool(vec[2] > 0.5),
|
|
"D": bool(vec[3] > 0.5),
|
|
}
|
|
if dim == 4:
|
|
return {
|
|
"W": bool(vec[0] > 0.5),
|
|
"S": bool(vec[1] > 0.5),
|
|
"A": bool(vec[2] > 0.5),
|
|
"D": bool(vec[3] > 0.5),
|
|
}
|
|
if dim == 2:
|
|
return {
|
|
"W": bool(vec[0] > 0.5),
|
|
"S": bool(vec[1] > 0.5),
|
|
"A": False,
|
|
"D": False,
|
|
}
|
|
return {"W": False, "S": False, "A": False, "D": False}
|
|
|
|
|
|
def draw_rounded_rectangle(
|
|
image: np.ndarray,
|
|
top_left: tuple[int, int],
|
|
bottom_right: tuple[int, int],
|
|
color: tuple[int, int, int],
|
|
radius: int = 10,
|
|
alpha: float = 0.5,
|
|
) -> None:
|
|
if not _require_cv2():
|
|
return
|
|
|
|
overlay = image.copy()
|
|
x1, y1 = top_left
|
|
x2, y2 = bottom_right
|
|
|
|
cv2.rectangle(overlay, (x1 + radius, y1), (x2 - radius, y2), color, -1)
|
|
cv2.rectangle(overlay, (x1, y1 + radius), (x2, y2 - radius), color, -1)
|
|
cv2.ellipse(overlay, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, -1)
|
|
cv2.ellipse(overlay, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, -1)
|
|
cv2.ellipse(overlay, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, -1)
|
|
cv2.ellipse(overlay, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, -1)
|
|
cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
|
|
|
|
|
|
def draw_keys_on_frame(
|
|
frame: np.ndarray,
|
|
keys: dict[str, bool],
|
|
key_size: tuple[int, int] = (30, 30),
|
|
top_margin: int = 15,
|
|
) -> None:
|
|
if not _require_cv2():
|
|
return
|
|
|
|
left_margin = 15
|
|
gap = 3
|
|
key_positions = {
|
|
"W": (left_margin + key_size[0] + gap, top_margin),
|
|
"A": (left_margin, top_margin + key_size[1] + gap),
|
|
"S": (left_margin + key_size[0] + gap, top_margin + key_size[1] + gap),
|
|
"D": (
|
|
left_margin + (key_size[0] + gap) * 2,
|
|
top_margin + key_size[1] + gap,
|
|
),
|
|
}
|
|
|
|
for key, (x, y) in key_positions.items():
|
|
pressed = keys.get(key, False)
|
|
color = (0, 255, 0) if pressed else (200, 200, 200)
|
|
alpha = 0.8 if pressed else 0.5
|
|
draw_rounded_rectangle(
|
|
frame,
|
|
(x, y),
|
|
(x + key_size[0], y + key_size[1]),
|
|
color,
|
|
radius=5,
|
|
alpha=alpha,
|
|
)
|
|
text_size = cv2.getTextSize(key, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
|
|
text_x = x + (key_size[0] - text_size[0]) // 2
|
|
text_y = y + (key_size[1] + text_size[1]) // 2
|
|
cv2.putText(
|
|
frame,
|
|
key,
|
|
(text_x, text_y),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.5,
|
|
(0, 0, 0),
|
|
1,
|
|
)
|
|
|
|
|
|
def draw_mouse_on_frame(
|
|
frame: np.ndarray,
|
|
yaw: float,
|
|
pitch: float,
|
|
top_margin: int = 15,
|
|
) -> None:
|
|
if not _require_cv2():
|
|
return
|
|
|
|
_, width, _ = frame.shape
|
|
right_margin = 15
|
|
crosshair_radius = 25
|
|
crosshair_x = width - right_margin - crosshair_radius
|
|
crosshair_y = top_margin + crosshair_radius
|
|
|
|
dx = int(yaw * crosshair_radius * 8)
|
|
dy = int(-pitch * crosshair_radius * 8)
|
|
max_arrow = crosshair_radius - 5
|
|
dx = max(-max_arrow, min(max_arrow, dx))
|
|
dy = max(-max_arrow, min(max_arrow, dy))
|
|
|
|
cv2.circle(frame, (crosshair_x, crosshair_y), crosshair_radius, (50, 50, 50), -1)
|
|
cv2.circle(frame, (crosshair_x, crosshair_y), crosshair_radius, (200, 200, 200), 1)
|
|
cv2.line(
|
|
frame,
|
|
(crosshair_x - crosshair_radius + 5, crosshair_y),
|
|
(crosshair_x + crosshair_radius - 5, crosshair_y),
|
|
(100, 100, 100),
|
|
1,
|
|
)
|
|
cv2.line(
|
|
frame,
|
|
(crosshair_x, crosshair_y - crosshair_radius + 5),
|
|
(crosshair_x, crosshair_y + crosshair_radius - 5),
|
|
(100, 100, 100),
|
|
1,
|
|
)
|
|
|
|
if abs(dx) > 1 or abs(dy) > 1:
|
|
cv2.arrowedLine(
|
|
frame,
|
|
(crosshair_x, crosshair_y),
|
|
(crosshair_x + dx, crosshair_y + dy),
|
|
(0, 255, 0),
|
|
2,
|
|
tipLength=0.3,
|
|
)
|
|
|
|
|
|
def overlay_validation_actions_on_frames(
|
|
frames: list[np.ndarray],
|
|
keyboard_cond: np.ndarray | torch.Tensor | None = None,
|
|
mouse_cond: np.ndarray | torch.Tensor | None = None,
|
|
) -> list[np.ndarray]:
|
|
if (keyboard_cond is None and mouse_cond is None) or not _require_cv2():
|
|
return frames
|
|
|
|
if keyboard_cond is not None and torch.is_tensor(keyboard_cond):
|
|
keyboard_cond = keyboard_cond.detach().cpu().float().numpy()
|
|
if mouse_cond is not None and torch.is_tensor(mouse_cond):
|
|
mouse_cond = mouse_cond.detach().cpu().float().numpy()
|
|
|
|
if keyboard_cond is not None:
|
|
keyboard_cond = np.asarray(keyboard_cond, dtype=np.float32)
|
|
if keyboard_cond.ndim == 3 and keyboard_cond.shape[0] == 1:
|
|
keyboard_cond = keyboard_cond[0]
|
|
if mouse_cond is not None:
|
|
mouse_cond = np.asarray(mouse_cond, dtype=np.float32)
|
|
if mouse_cond.ndim == 3 and mouse_cond.shape[0] == 1:
|
|
mouse_cond = mouse_cond[0]
|
|
|
|
processed_frames: list[np.ndarray] = []
|
|
for frame_idx, frame in enumerate(frames):
|
|
frame = np.ascontiguousarray(frame.copy())
|
|
if keyboard_cond is not None and frame_idx < len(keyboard_cond):
|
|
draw_keys_on_frame(
|
|
frame,
|
|
_matrixgame2_overlay_keys_from_keyboard(keyboard_cond[frame_idx]),
|
|
)
|
|
if mouse_cond is not None and frame_idx < len(mouse_cond):
|
|
# MG action convention: mouse[0]=pitch, mouse[1]=yaw.
|
|
pitch = float(mouse_cond[frame_idx, 0])
|
|
yaw = float(mouse_cond[frame_idx, 1])
|
|
draw_mouse_on_frame(frame, yaw=yaw, pitch=pitch)
|
|
processed_frames.append(frame)
|
|
|
|
return processed_frames
|
|
|
|
|
|
def scale_keyboard_condition(
|
|
keyboard_condition: torch.Tensor,
|
|
scale: float,
|
|
) -> torch.Tensor:
|
|
if float(scale) == 1.0:
|
|
return keyboard_condition
|
|
return keyboard_condition * float(scale)
|
|
|
|
|
|
def expand_action_to_frames(action: dict, num_frames: int) -> tuple[torch.Tensor, torch.Tensor]:
|
|
result = {}
|
|
for key, tensor in action.items():
|
|
if tensor is not None:
|
|
# Expand to [num_frames, D] then unsqueeze to [1, num_frames, D]
|
|
result[key] = tensor.unsqueeze(0).repeat(num_frames, 1).unsqueeze(0)
|
|
else:
|
|
result[key] = None
|
|
|
|
if "mouse" not in result or result["mouse"] is None:
|
|
# keyboard device if available, otherwise default
|
|
device = result.get("keyboard", torch.tensor(
|
|
[])).device if result.get("keyboard") is not None else get_local_torch_device()
|
|
result["mouse"] = torch.zeros(1, num_frames, 2, device=device)
|
|
|
|
return result["keyboard"], result["mouse"]
|
|
|
|
|
|
def get_current_action(mode="universal"):
|
|
|
|
CAM_VALUE = 0.1
|
|
if mode == 'universal':
|
|
logger.info("")
|
|
logger.info('-' * 30)
|
|
logger.info("PRESS [I, K, J, L, U] FOR CAMERA TRANSFORM\n (I: up, K: down, J: left, L: right, U: no move)")
|
|
logger.info("PRESS [W, S, A, D, Q] FOR MOVEMENT\n (W: forward, S: back, A: left, D: right, Q: no move)")
|
|
logger.info('-' * 30)
|
|
CAMERA_VALUE_MAP = {
|
|
"i": [CAM_VALUE, 0],
|
|
"k": [-CAM_VALUE, 0],
|
|
"j": [0, -CAM_VALUE],
|
|
"l": [0, CAM_VALUE],
|
|
"u": [0, 0]
|
|
}
|
|
KEYBOARD_IDX = {"w": [1, 0, 0, 0], "s": [0, 1, 0, 0], "a": [0, 0, 1, 0], "d": [0, 0, 0, 1], "q": [0, 0, 0, 0]}
|
|
flag = 0
|
|
while flag != 1:
|
|
try:
|
|
idx_mouse = input('Please input the mouse action (e.g. `U`):\n').strip().lower()
|
|
idx_keyboard = input('Please input the keyboard action (e.g. `W`):\n').strip().lower()
|
|
if idx_mouse in CAMERA_VALUE_MAP and idx_keyboard in KEYBOARD_IDX:
|
|
flag = 1
|
|
except Exception:
|
|
pass
|
|
mouse_cond = torch.tensor(CAMERA_VALUE_MAP[idx_mouse]).cuda()
|
|
keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard]).cuda()
|
|
elif mode == 'gta_drive':
|
|
logger.info("")
|
|
logger.info('-' * 30)
|
|
logger.info("PRESS [W, S, A, D, Q] FOR MOVEMENT\n (W: forward, S: back, A: left, D: right, Q: no move)")
|
|
logger.info('-' * 30)
|
|
CAMERA_VALUE_MAP = {"a": [0, -CAM_VALUE], "d": [0, CAM_VALUE], "q": [0, 0]}
|
|
KEYBOARD_IDX = {"w": [1, 0], "s": [0, 1], "q": [0, 0]}
|
|
flag = 0
|
|
while flag != 1:
|
|
try:
|
|
indexes = input(
|
|
'Please input the actions (split with ` `):\n(e.g. `W` for forward, `W A` for forward and left)\n'
|
|
).strip().lower().split(' ')
|
|
idx_mouse = []
|
|
idx_keyboard = []
|
|
for i in indexes:
|
|
if i in CAMERA_VALUE_MAP.keys():
|
|
idx_mouse += [i]
|
|
elif i in KEYBOARD_IDX.keys():
|
|
idx_keyboard += [i]
|
|
if len(idx_mouse) == 0:
|
|
idx_mouse += ['q']
|
|
if len(idx_keyboard) == 0:
|
|
idx_keyboard += ['q']
|
|
assert idx_mouse in [['a'], ['d'], ['q']] and idx_keyboard in [['q'], ['w'], ['s']]
|
|
flag = 1
|
|
except Exception:
|
|
pass
|
|
mouse_cond = torch.tensor(CAMERA_VALUE_MAP[idx_mouse[0]]).cuda()
|
|
keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard[0]]).cuda()
|
|
elif mode == 'templerun':
|
|
logger.info("")
|
|
logger.info('-' * 30)
|
|
logger.info(
|
|
"PRESS [W, S, A, D, Z, C, Q] FOR ACTIONS\n (W: jump, S: slide, A: left side, D: right side, Z: turn left, C: turn right, Q: no move)"
|
|
)
|
|
logger.info('-' * 30)
|
|
KEYBOARD_IDX = {
|
|
"w": [0, 1, 0, 0, 0, 0, 0],
|
|
"s": [0, 0, 1, 0, 0, 0, 0],
|
|
"a": [0, 0, 0, 0, 0, 1, 0],
|
|
"d": [0, 0, 0, 0, 0, 0, 1],
|
|
"z": [0, 0, 0, 1, 0, 0, 0],
|
|
"c": [0, 0, 0, 0, 1, 0, 0],
|
|
"q": [1, 0, 0, 0, 0, 0, 0]
|
|
}
|
|
flag = 0
|
|
while flag != 1:
|
|
try:
|
|
idx_keyboard = input(
|
|
'Please input the action: \n(e.g. `W` for forward, `Z` for turning left)\n').strip().lower()
|
|
if idx_keyboard in KEYBOARD_IDX.keys():
|
|
flag = 1
|
|
except Exception:
|
|
pass
|
|
keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard]).cuda()
|
|
|
|
if mode != 'templerun':
|
|
return {"mouse": mouse_cond, "keyboard": keyboard_cond}
|
|
return {"keyboard": keyboard_cond}
|
|
|
|
|
|
async def get_current_action_async(mode="universal"):
|
|
return await asyncio.to_thread(get_current_action, mode)
|
|
|
|
|
|
def load_initial_image(image_path: str = None) -> Image.Image:
|
|
if image_path and os.path.exists(image_path):
|
|
return Image.open(image_path).convert("RGB")
|
|
logger.warning("No image provided, creating placeholder...")
|
|
return Image.new("RGB", (640, 352), (128, 128, 128))
|
|
|
|
|
|
def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None):
|
|
if keyboard_dim not in (2, 4, 6, 7):
|
|
raise ValueError(f"keyboard_dim must be 2, 4, 6, or 7, got {keyboard_dim}")
|
|
if num_frames % 4 != 1:
|
|
raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.")
|
|
|
|
# Set seed for reproducibility if provided
|
|
if seed is not None:
|
|
random.seed(seed)
|
|
|
|
num_samples_per_action = 4
|
|
|
|
# Define actions based on keyboard_dim
|
|
if keyboard_dim == 4:
|
|
# Universal model: W, S, A, D
|
|
actions_single_action = ["forward", "left", "right"]
|
|
actions_double_action = ["forward_left", "forward_right"]
|
|
actions_single_camera = ["camera_l", "camera_r"]
|
|
keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3}
|
|
elif keyboard_dim == 2:
|
|
# GTA model: W, S only (steering via mouse)
|
|
actions_single_action = ["forward", "back"]
|
|
actions_double_action = []
|
|
actions_single_camera = ["camera_l", "camera_r"]
|
|
keyboard_idx = {"forward": 0, "back": 1}
|
|
elif keyboard_dim == 6:
|
|
actions_single_action = ["forward", "back", "left", "right"]
|
|
actions_double_action = ["forward_left", "forward_right"]
|
|
actions_single_camera = ["camera_l", "camera_r"]
|
|
keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3, "t1": 4, "t2": 5}
|
|
else: # keyboard_dim == 7
|
|
# Temple Run model: still, w, s, left, right, a, d (no mouse)
|
|
actions_single_action = ["forward", "back", "left", "right"]
|
|
actions_double_action = []
|
|
actions_single_camera = [] # No mouse for Temple Run
|
|
keyboard_idx = {"still": 0, "forward": 1, "back": 2, "left": 3, "right": 4, "a": 5, "d": 6}
|
|
|
|
actions_to_test = (actions_double_action * 5 + actions_single_camera * 5 + actions_single_action * 5)
|
|
for action in (actions_single_action + actions_double_action):
|
|
for camera in actions_single_camera:
|
|
actions_to_test.append(f"{action}_{camera}")
|
|
|
|
# Ensure we have at least some actions
|
|
if not actions_to_test:
|
|
actions_to_test = actions_single_action * 5
|
|
|
|
base_action = actions_single_action + actions_single_camera
|
|
cam_value = 0.1
|
|
camera_value_map = {
|
|
"camera_up": [cam_value, 0],
|
|
"camera_down": [-cam_value, 0],
|
|
"camera_l": [0, -cam_value],
|
|
"camera_r": [0, cam_value],
|
|
"camera_ur": [cam_value, cam_value],
|
|
"camera_ul": [cam_value, -cam_value],
|
|
"camera_dr": [-cam_value, cam_value],
|
|
"camera_dl": [-cam_value, -cam_value],
|
|
}
|
|
|
|
data = []
|
|
for action_name in actions_to_test:
|
|
keyboard_condition = torch.zeros((num_samples_per_action, keyboard_dim))
|
|
mouse_condition = torch.zeros((num_samples_per_action, 2))
|
|
|
|
for sub_act in base_action:
|
|
if sub_act not in action_name:
|
|
continue
|
|
if sub_act in camera_value_map:
|
|
mouse_condition = torch.tensor(
|
|
[camera_value_map[sub_act] for _ in range(num_samples_per_action)],
|
|
dtype=mouse_condition.dtype,
|
|
)
|
|
elif sub_act in keyboard_idx:
|
|
keyboard_condition[:, keyboard_idx[sub_act]] = 1
|
|
|
|
data.append({
|
|
"keyboard_condition": keyboard_condition,
|
|
"mouse_condition": mouse_condition,
|
|
})
|
|
|
|
keyboard_condition = torch.zeros((num_frames, keyboard_dim))
|
|
mouse_condition = torch.zeros((num_frames, 2))
|
|
current_frame = 0
|
|
selections = [12]
|
|
|
|
while current_frame < num_frames:
|
|
rd_frame = selections[random.randint(0, len(selections) - 1)]
|
|
entry = data[random.randint(0, len(data) - 1)]
|
|
key_seq = entry["keyboard_condition"]
|
|
mouse_seq = entry["mouse_condition"]
|
|
|
|
if current_frame == 0:
|
|
keyboard_condition[:1] = key_seq[:1]
|
|
mouse_condition[:1] = mouse_seq[:1]
|
|
current_frame = 1
|
|
else:
|
|
rd_frame = min(rd_frame, num_frames - current_frame)
|
|
repeat_time = rd_frame // 4
|
|
keyboard_condition[current_frame:current_frame + rd_frame] = key_seq.repeat(repeat_time, 1)
|
|
mouse_condition[current_frame:current_frame + rd_frame] = mouse_seq.repeat(repeat_time, 1)
|
|
current_frame += rd_frame
|
|
|
|
return {"keyboard": keyboard_condition, "mouse": mouse_condition}
|
|
|
|
|
|
def parse_config(config, mode="universal"):
|
|
assert mode in ['universal', 'gta_drive', 'templerun']
|
|
key_data = {}
|
|
mouse_data = {}
|
|
if mode != 'templerun':
|
|
key, mouse = config
|
|
else:
|
|
key = config
|
|
|
|
for i in range(len(key)):
|
|
if mode == 'templerun':
|
|
still, w, s, left, right, a, d = key[i]
|
|
elif mode == 'universal':
|
|
w, s, a, d = key[i]
|
|
else:
|
|
w, s, a, d = key[i][0], key[i][1], mouse[i][1] < 0, mouse[i][1] > 0
|
|
if mode == 'universal':
|
|
mouse_y, mouse_x = mouse[i]
|
|
mouse_y = -1 * mouse_y
|
|
|
|
key_data[i] = {"W": bool(w), "A": bool(a), "S": bool(s), "D": bool(d)}
|
|
if mode == 'templerun':
|
|
key_data[i].update({"left": bool(left), "right": bool(right)})
|
|
|
|
if mode == 'universal':
|
|
if i == 0:
|
|
mouse_data[i] = (320, 352 // 2)
|
|
else:
|
|
global_scale_factor = 0.1
|
|
mouse_scale_x = 15 * global_scale_factor
|
|
mouse_scale_y = 15 * 4 * global_scale_factor
|
|
mouse_data[i] = (
|
|
mouse_data[i - 1][0] + mouse_x * mouse_scale_x,
|
|
mouse_data[i - 1][1] + mouse_y * mouse_scale_y,
|
|
)
|
|
return key_data, mouse_data
|
|
|
|
|
|
# NOTE: drawing functions are commented out to avoid cv2/libGL dependency.
|
|
#
|
|
# def draw_rounded_rectangle(image, top_left, bottom_right, color, radius=10, alpha=0.5):
|
|
# overlay = image.copy()
|
|
# x1, y1 = top_left
|
|
# x2, y2 = bottom_right
|
|
#
|
|
# cv2.rectangle(overlay, (x1 + radius, y1), (x2 - radius, y2), color, -1)
|
|
# cv2.rectangle(overlay, (x1, y1 + radius), (x2, y2 - radius), color, -1)
|
|
# cv2.ellipse(overlay, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, -1)
|
|
# cv2.ellipse(overlay, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, -1)
|
|
# cv2.ellipse(overlay, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, -1)
|
|
# cv2.ellipse(overlay, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, -1)
|
|
# cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image)
|
|
#
|
|
# def draw_keys_on_frame(frame, keys, key_size=(80, 50), spacing=20, bottom_margin=30, mode='universal'):
|
|
# h, w, _ = frame.shape
|
|
# horison_shift = 90
|
|
# vertical_shift = -20
|
|
# horizon_shift_all = 50
|
|
# key_positions = {
|
|
# "W": (w // 2 - key_size[0] // 2 - horison_shift - horizon_shift_all,
|
|
# h - bottom_margin - key_size[1] * 2 + vertical_shift - 20),
|
|
# "A": (w // 2 - key_size[0] * 2 + 5 - horison_shift - horizon_shift_all,
|
|
# h - bottom_margin - key_size[1] + vertical_shift),
|
|
# "S": (w // 2 - key_size[0] // 2 - horison_shift - horizon_shift_all,
|
|
# h - bottom_margin - key_size[1] + vertical_shift),
|
|
# "D": (w // 2 + key_size[0] - 5 - horison_shift - horizon_shift_all,
|
|
# h - bottom_margin - key_size[1] + vertical_shift),
|
|
# }
|
|
# key_icon = {"W": "W", "A": "A", "S": "S", "D": "D", "left": "left", "right": "right"}
|
|
# if mode == 'templerun':
|
|
# key_positions.update({
|
|
# "left": (w // 2 + key_size[0] * 2 + spacing * 2 - horison_shift - horizon_shift_all,
|
|
# h - bottom_margin - key_size[1] + vertical_shift),
|
|
# "right": (w // 2 + key_size[0] * 3 + spacing * 7 - horison_shift - horizon_shift_all,
|
|
# h - bottom_margin - key_size[1] + vertical_shift)
|
|
# })
|
|
#
|
|
# for key, (x, y) in key_positions.items():
|
|
# is_pressed = keys.get(key, False)
|
|
# top_left = (x, y)
|
|
# if key in ["left", "right"]:
|
|
# bottom_right = (x + key_size[0] + 40, y + key_size[1])
|
|
# else:
|
|
# bottom_right = (x + key_size[0], y + key_size[1])
|
|
#
|
|
# color = (0, 255, 0) if is_pressed else (200, 200, 200)
|
|
# alpha = 0.8 if is_pressed else 0.5
|
|
# draw_rounded_rectangle(frame, top_left, bottom_right, color, radius=10, alpha=alpha)
|
|
#
|
|
# text_size = cv2.getTextSize(key, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)[0]
|
|
# if key in ["left", "right"]:
|
|
# text_x = x + (key_size[0] + 40 - text_size[0]) // 2
|
|
# else:
|
|
# text_x = x + (key_size[0] - text_size[0]) // 2
|
|
# text_y = y + (key_size[1] + text_size[1]) // 2
|
|
# cv2.putText(frame, key_icon[key], (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)
|
|
#
|
|
# def overlay_icon(frame, icon, position, scale=1.0, rotation=0):
|
|
# x, y = position
|
|
# h, w, _ = icon.shape
|
|
#
|
|
# scaled_width = int(w * scale)
|
|
# scaled_height = int(h * scale)
|
|
# icon_resized = cv2.resize(icon, (scaled_width, scaled_height), interpolation=cv2.INTER_AREA)
|
|
#
|
|
# center = (scaled_width // 2, scaled_height // 2)
|
|
# rotation_matrix = cv2.getRotationMatrix2D(center, rotation, 1.0)
|
|
# icon_rotated = cv2.warpAffine(
|
|
# icon_resized, rotation_matrix, (scaled_width, scaled_height),
|
|
# flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0)
|
|
# )
|
|
#
|
|
# h, w, _ = icon_rotated.shape
|
|
# frame_h, frame_w, _ = frame.shape
|
|
#
|
|
# top_left_x = max(0, int(x - w // 2))
|
|
# top_left_y = max(0, int(y - h // 2))
|
|
# bottom_right_x = min(frame_w, int(x + w // 2))
|
|
# bottom_right_y = min(frame_h, int(y + h // 2))
|
|
#
|
|
# icon_x_start = max(0, int(-x + w // 2))
|
|
# icon_y_start = max(0, int(-y + h // 2))
|
|
# icon_x_end = icon_x_start + (bottom_right_x - top_left_x)
|
|
# icon_y_end = icon_y_start + (bottom_right_y - top_left_y)
|
|
#
|
|
# icon_region = icon_rotated[icon_y_start:icon_y_end, icon_x_start:icon_x_end]
|
|
# alpha = icon_region[:, :, 3] / 255.0
|
|
# icon_rgb = icon_region[:, :, :3]
|
|
#
|
|
# frame_region = frame[top_left_y:bottom_right_y, top_left_x:bottom_right_x]
|
|
# for c in range(3):
|
|
# frame_region[:, :, c] = (1 - alpha) * frame_region[:, :, c] + alpha * icon_rgb[:, :, c]
|
|
# frame[top_left_y:bottom_right_y, top_left_x:bottom_right_x] = frame_region
|
|
#
|
|
# def process_video(input_video, output_video, config, mouse_icon_path,
|
|
# mouse_scale=1.0, mouse_rotation=0, process_icon=True, mode='universal'):
|
|
# key_data, mouse_data = parse_config(config, mode=mode)
|
|
# fps = 12
|
|
#
|
|
# mouse_icon = cv2.imread(mouse_icon_path, cv2.IMREAD_UNCHANGED)
|
|
#
|
|
# out_video = []
|
|
# for frame_idx, frame in enumerate(input_video):
|
|
# frame = np.ascontiguousarray(frame)
|
|
# if process_icon:
|
|
# keys = key_data.get(frame_idx, {"W": False, "A": False, "S": False, "D": False, "left": False, "right": False})
|
|
# draw_keys_on_frame(frame, keys, key_size=(50, 50), spacing=10, bottom_margin=20, mode=mode)
|
|
# if mode == 'universal':
|
|
# frame_width = frame.shape[1]
|
|
# frame_height = frame.shape[0]
|
|
# mouse_position = mouse_data.get(frame_idx, (frame_width // 2, frame_height // 2))
|
|
# overlay_icon(frame, mouse_icon, mouse_position, scale=mouse_scale, rotation=mouse_rotation)
|
|
# out_video.append(frame / 255)
|
|
#
|
|
# export_to_video(out_video, output_video, fps=fps)
|
|
# logger.info(f"Video saved to {output_video}")
|