File size: 1,501 Bytes
cb91e74 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import os
import cv2
import numpy as np
from pathlib import Path
def load_templates(cursor_dir):
templates = []
for fname in os.listdir(cursor_dir):
if fname.endswith('.png'):
templates.append(cv2.imread(os.path.join(cursor_dir, fname), cv2.IMREAD_UNCHANGED))
return templates
def track_cursor(frames_dir, cursor_dir, output_json):
import json
templates = load_templates(cursor_dir)
results = []
for frame_file in sorted(os.listdir(frames_dir)):
if not frame_file.endswith('.png'):
continue
frame_path = os.path.join(frames_dir, frame_file)
frame = cv2.imread(frame_path)
best_match = None
best_val = -np.inf
for template in templates:
res = cv2.matchTemplate(frame, template[..., :3], cv2.TM_CCOEFF_NORMED)
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
if max_val > best_val:
best_val = max_val
best_match = (max_loc, template.shape[:2], max_val)
if best_match:
(x, y), (h, w), score = best_match
results.append({"frame": frame_file, "cursor_pos": [int(x), int(y)], "score": float(score)})
with open(output_json, "w") as f:
json.dump(results, f, indent=2)
print(f"Cursor tracking complete. Results saved to {output_json}")
if __name__ == "__main__":
track_cursor("frames", "cursors", "annotations/cursor_positions.json")
|