Skip to content

Commit

Permalink
feat: Add RecordingWidget for video recording functionality
Browse files Browse the repository at this point in the history
- Created RecordingWidget class to handle video recording from canvas
- Added methods to start, stop, and capture frames during recording
- Integrated folder selection for saving recordings
- Ensured frames are converted to BGR format before saving
- Added debug messages for recording dimensions and frame shape
- Included record button in the toolbar with toggle functionality
  • Loading branch information
healthonrails committed Dec 20, 2024
1 parent 64278a0 commit ebf1013
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 1 deletion.
4 changes: 4 additions & 0 deletions annolid/gui/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
from annolid.gui.widgets.downsample_videos_dialog import VideoRescaleWidget
from annolid.gui.widgets.convert_sleap_dialog import ConvertSleapDialog
from annolid.gui.widgets.extract_keypoints_dialog import ExtractShapeKeyPointsDialog
from annolid.gui.widgets import RecordingWidget
from annolid.gui.widgets.convert_labelme2csv_dialog import LabelmeJsonToCsvDialog
from annolid.postprocessing.quality_control import pred_dict_to_labelme
from annolid.annotation.timestamps import convert_frame_number_to_time
Expand Down Expand Up @@ -479,6 +480,8 @@ def __init__(self,
self._grounding_sam
)

self.recording_widget = RecordingWidget(self.canvas)

_action_tools = list(self.actions.tool)
_action_tools.insert(0, frames)
_action_tools.insert(1, open_video)
Expand All @@ -493,6 +496,7 @@ def __init__(self,
_action_tools.append(quality_control)
_action_tools.append(colab)
_action_tools.append(visualization)
_action_tools.append(self.recording_widget.record_action)

self.actions.tool = tuple(_action_tools)
self.tools.clear()
Expand Down
3 changes: 2 additions & 1 deletion annolid/gui/widgets/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,5 @@
from annolid.gui.widgets.progressing_dialog import ProgressingWindow
from annolid.gui.widgets.quality_control_dialog import QualityControlDialog
from annolid.gui.widgets.about_dialog import SystemInfoDialog
from annolid.gui.widgets.flags import FlagTableWidget
from annolid.gui.widgets.flags import FlagTableWidget
from annolid.gui.widgets.video_recording import RecordingWidget
107 changes: 107 additions & 0 deletions annolid/gui/widgets/video_recording.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
from qtpy import QtWidgets, QtGui, QtCore
import cv2
from datetime import datetime
from labelme import utils
from pathlib import Path
import numpy as np


class RecordingWidget(QtWidgets.QWidget):
def __init__(self, canvas, fps=30, parent=None):
super().__init__(parent)
self.canvas = canvas
self.fps = fps
self.is_recording = False
self.video_writer = None
self.output_filename = None
self.capture_timer = None
self.here = Path(__file__).resolve().parent.parent
self.record_icon = QtGui.QIcon(str(self.here / "icons/record.png"))
self.stop_record_icon = QtGui.QIcon(
str(self.here / "icons/stop_record.png"))

# Add a record button to the toolbar
self.record_action = QtWidgets.QAction(
self.record_icon, "Record", self)
self.record_action.setCheckable(True)
self.record_action.toggled.connect(self.toggle_record)

def toggle_record(self, checked):
if checked:
self.start_recording()
else:
self.stop_recording()

def start_recording(self):
# Prompt the user to select a folder
folder = QtWidgets.QFileDialog.getExistingDirectory(
self, "Select Folder")
if not folder:
QtWidgets.QMessageBox.warning(
self, "No Folder Selected", "Please select a folder to save the recording.")
self.record_action.setChecked(False)
return

self.is_recording = True
self.record_action.setIcon(self.stop_record_icon)
self.record_action.setText("Stop Recording")

# Get the canvas pixmap dimensions
pixmap = self.canvas.grab()
width = pixmap.width()
height = pixmap.height()
fps = self.fps if self.fps is not None else 30 # Default to 30 if not available

# Define the codec and create VideoWriter object
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Use 'mp4v' codec for .mp4
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
# Store the filename with selected folder
self.output_filename = f"{folder}/canvas_recording_{timestamp}.mp4"
self.video_writer = cv2.VideoWriter(
self.output_filename, fourcc, fps, (width, height))

if not self.video_writer.isOpened():
QtWidgets.QMessageBox.critical(
self, "Error", f"Could not open video writer for {self.output_filename}")
self.is_recording = False
self.record_action.setChecked(False)
self.record_action.setIcon(self.record_icon)
self.record_action.setText("Record")
self.video_writer = None
return

# Start a timer to capture frames
self.capture_timer = QtCore.QTimer(self)
self.capture_timer.timeout.connect(self.capture_frame)
self.capture_timer.start(int(1000 / fps)) # Interval in milliseconds

def capture_frame(self):
"""Capture the current frame from the canvas."""
if self.is_recording:
# Get the current pixmap from the canvas
pixmap = self.canvas.grab()
qimage = pixmap.toImage()

# Convert QImage to numpy array
frame = utils.img_qt_to_arr(qimage)
frame_bgr = frame[:, :, :3]
# Write the frame to the video writer
self.video_writer.write(frame_bgr)

def stop_recording(self):
self.is_recording = False
self.record_action.setIcon(self.stop_record_icon)
self.record_action.setText("Record")

if self.capture_timer and self.capture_timer.isActive():
self.capture_timer.stop()
self.capture_timer.deleteLater()
self.capture_timer = None

if self.video_writer and self.video_writer.isOpened():
self.video_writer.release()
self.video_writer = None

# Display message with the saved filename
QtWidgets.QMessageBox.information(
self, "Recording Saved", f"Canvas recording saved to {self.output_filename}")

0 comments on commit ebf1013

Please sign in to comment.