Webカメラのストリーム映像から角度計測をするスクリプト

import sys
import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button
import time

plt.rcParams['font.family'] = 'MS Gothic'

cap = cv2.VideoCapture(0)
if not cap.isOpened():
    print("error")
    sys.exit()

cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

fig, ax = plt.subplots(figsize=(10, 8))
ax.set_xlim(0, frame_width)
ax.set_ylim(frame_height, 0)
ax.set_title("WebCam Stream ")
ax.axis('off')

ret, frame = cap.read()
if not ret:
    print("error。")
    sys.exit()
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
im = ax.imshow(frame_rgb)

info_text = ax.text(
    0.05, 0.95, '', transform=ax.transAxes, fontsize=10,
    verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5)
)

init_points = [
    (frame_width * 0.3, frame_height * 0.3),  
    (frame_width * 0.7, frame_height * 0.3),  
    (frame_width * 0.3, frame_height * 0.7),  
    (frame_width * 0.7, frame_height * 0.7),  
]

points = [ax.plot(x, y, 'wo', markersize=8, picker=5)[0] for x, y in init_points]
lines = [
    ax.plot([init_points[0][0], init_points[1][0]], [init_points[0][1], init_points[1][1]], 'w-', lw=2)[0],
    ax.plot([init_points[2][0], init_points[3][0]], [init_points[2][1], init_points[3][1]], 'w-', lw=2)[0]
]

def angle_between_lines(p1, p2, p3, p4):
    vec1 = np.array([p2[0] - p1[0], p2[1] - p1[1]])
    vec2 = np.array([p4[0] - p3[0], p4[1] - p3[1]])
    unit_vec1 = vec1 / np.linalg.norm(vec1)
    unit_vec2 = vec2 / np.linalg.norm(vec2)
    dot_product = np.clip(np.dot(unit_vec1, unit_vec2), -1.0, 1.0)
    angle_rad = np.arccos(dot_product)
    angle_deg = np.degrees(angle_rad)
    return angle_deg

def update_info():
    coords = [ (p.get_data()[0][0], p.get_data()[1][0]) for p in points ]
    lines[0].set_data([coords[0][0], coords[1][0]], [coords[0][1], coords[1][1]])
    lines[1].set_data([coords[2][0], coords[3][0]], [coords[2][1], coords[3][1]])
    angle = angle_between_lines(coords[0], coords[1], coords[2], coords[3])
    info_text.set_text(f"angle: {angle:.2f}°")
    fig.canvas.draw_idle()

class DraggablePoint:
    def __init__(self, point, update_callback=None):
        self.point = point
        self.press = None
        self.update_callback = update_callback

    def connect(self):
        self.cidpress = self.point.figure.canvas.mpl_connect('button_press_event', self.on_press)
        self.cidrelease = self.point.figure.canvas.mpl_connect('button_release_event', self.on_release)
        self.cidmotion = self.point.figure.canvas.mpl_connect('motion_notify_event', self.on_motion)

    def on_press(self, event):
        if event.inaxes != self.point.axes:
            return
        contains, _ = self.point.contains(event)
        if not contains:
            return
        x0, y0 = self.point.get_data()
        self.press = (x0[0], y0[0], event.xdata, event.ydata)

    def on_motion(self, event):
        if self.press is None or event.inaxes != self.point.axes:
            return
        x0, y0, xpress, ypress = self.press
        dx = event.xdata - xpress
        dy = event.ydata - ypress
        new_x = x0 + dx
        new_y = y0 + dy
        self.point.set_data([new_x], [new_y])
        if self.update_callback:
            self.update_callback()
        self.point.figure.canvas.draw_idle()

    def on_release(self, event):
        self.press = None
        self.point.figure.canvas.draw_idle()

draggables = []
for pt in points:
    dp = DraggablePoint(pt, update_callback=update_info)
    dp.connect()
    draggables.append(dp)
update_info()

# StopとSnapshotボタンを追加
ax_stop = fig.add_axes([0.1, 0.92, 0.1, 0.05])
ax_snapshot = fig.add_axes([0.25, 0.92, 0.1, 0.05])
btn_stop = Button(ax_stop, 'Stop')
btn_snapshot = Button(ax_snapshot, 'Snapshot')

def stop_callback(event):
    global running
    running = False
    plt.close(fig)

def snapshot_callback(event):
    filename = f"snapshot_{int(time.time())}.png"
    fig.savefig(filename)
    print(f"Snapshot saved as {filename}")

btn_stop.on_clicked(stop_callback)
btn_snapshot.on_clicked(snapshot_callback)

running = True
while running and plt.fignum_exists(fig.number):
    ret, frame = cap.read()
    if ret:
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        im.set_data(frame_rgb)
    else:
        print("Failed to read frame")
    fig.canvas.draw()
    plt.pause(0.03)

cap.release()