draft video cap
This commit is contained in:
RyanOnTheInside
@@ -55,47 +55,143 @@ class FastWebcamCapture:
|
||||
return {
|
||||
"required": {
|
||||
"image": ("WEBCAM", {}),
|
||||
"capture_on_queue": ("BOOLEAN", {"default": True, "tooltip": "Whether to capture the frame when the node is queued"}),
|
||||
"record_seconds": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 60.0,
|
||||
"step": 0.5,
|
||||
"tooltip": "Length of video to record in seconds. 0 means capture a single image (no video)."
|
||||
}
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_TYPES = ("IMAGE", "AUDIO")
|
||||
RETURN_NAMES = ("image", "audio")
|
||||
FUNCTION = "process_capture"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def process_capture(self, image, capture_on_queue):
|
||||
def process_capture(self, image, record_seconds):
|
||||
"""
|
||||
Process a webcam image at native resolution.
|
||||
Process a webcam image or video at native resolution.
|
||||
|
||||
This node simply captures the webcam feed at its native resolution.
|
||||
This node captures the webcam feed at its native resolution.
|
||||
If record_seconds > 0, it will record a video for the specified duration.
|
||||
For resizing or cropping, use standard ComfyUI nodes after this one.
|
||||
"""
|
||||
# Check if we got a data URL
|
||||
if isinstance(image, str) and image.startswith("data:image/"):
|
||||
# Extract the base64 data after the comma
|
||||
base64_data = re.sub("^data:image/.+;base64,", "", image)
|
||||
import torchaudio
|
||||
import torch
|
||||
import numpy as np
|
||||
import time
|
||||
import os
|
||||
import tempfile
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
# Create empty audio dict as default
|
||||
empty_audio = {"waveform": torch.zeros((1, 16000)), "sample_rate": 16000}
|
||||
|
||||
# Check if we're receiving image or video data
|
||||
if isinstance(image, str):
|
||||
# Image data URL
|
||||
if image.startswith("data:image/"):
|
||||
# Extract the base64 data after the comma
|
||||
base64_data = re.sub("^data:image/.+;base64,", "", image)
|
||||
|
||||
# Convert base64 to PIL Image
|
||||
buffer = BytesIO(base64.b64decode(base64_data))
|
||||
pil_image = Image.open(buffer).convert("RGB")
|
||||
|
||||
# Log the image size we're receiving from the webcam
|
||||
print(f"Webcam image size: {pil_image.width}x{pil_image.height}")
|
||||
# Convert base64 to PIL Image
|
||||
buffer = BytesIO(base64.b64decode(base64_data))
|
||||
pil_image = Image.open(buffer).convert("RGB")
|
||||
|
||||
# Log the image size we're receiving from the webcam
|
||||
print(f"Webcam image size: {pil_image.width}x{pil_image.height}")
|
||||
|
||||
# Convert PIL to numpy array
|
||||
image = np.array(pil_image)
|
||||
# Convert PIL to numpy array
|
||||
image_np = np.array(pil_image)
|
||||
|
||||
# Convert to float32 and normalize to 0-1 range
|
||||
image = image.astype(np.float32) / 255.0
|
||||
# Convert to float32 and normalize to 0-1 range
|
||||
image_np = image_np.astype(np.float32) / 255.0
|
||||
|
||||
# Add batch dimension and convert to torch tensor
|
||||
# ComfyUI expects BHWC format
|
||||
image = torch.from_numpy(image)[None, ...]
|
||||
|
||||
return (image,)
|
||||
# Add batch dimension and convert to torch tensor
|
||||
# ComfyUI expects BHWC format
|
||||
image_tensor = torch.from_numpy(image_np)[None, ...]
|
||||
|
||||
return (image_tensor, empty_audio)
|
||||
|
||||
# Video data URL
|
||||
elif image.startswith("data:video/"):
|
||||
print("Received video data from webcam")
|
||||
|
||||
try:
|
||||
# Extract the base64 data after the comma
|
||||
mime_pattern = "data:video/.+;base64,"
|
||||
base64_data = re.sub(f"^{mime_pattern}", "", image)
|
||||
|
||||
# Create a temporary file to store the video
|
||||
with tempfile.NamedTemporaryFile(suffix='.webm', delete=False) as video_file:
|
||||
video_path = video_file.name
|
||||
video_file.write(base64.b64decode(base64_data))
|
||||
|
||||
# Process video with ffmpeg to extract frames
|
||||
frames_dir = tempfile.mkdtemp()
|
||||
frames_path = os.path.join(frames_dir, "frame_%04d.png")
|
||||
|
||||
# Extract frames with ffmpeg
|
||||
subprocess.run([
|
||||
"ffmpeg", "-i", video_path,
|
||||
"-vf", "fps=30",
|
||||
frames_path
|
||||
], check=True, capture_output=True)
|
||||
|
||||
# Load extracted frames
|
||||
frames = []
|
||||
frame_files = sorted([f for f in os.listdir(frames_dir) if f.startswith("frame_")])
|
||||
|
||||
if not frame_files:
|
||||
raise ValueError("No frames extracted from video")
|
||||
|
||||
# Load all frames
|
||||
for frame_file in frame_files:
|
||||
frame_path = os.path.join(frames_dir, frame_file)
|
||||
pil_frame = Image.open(frame_path).convert("RGB")
|
||||
frame_np = np.array(pil_frame).astype(np.float32) / 255.0
|
||||
frames.append(torch.from_numpy(frame_np))
|
||||
|
||||
# Extract audio with torchaudio
|
||||
try:
|
||||
waveform, sample_rate = torchaudio.load(video_path)
|
||||
audio = {"waveform": waveform, "sample_rate": sample_rate}
|
||||
except Exception as e:
|
||||
print(f"Error extracting audio: {e}")
|
||||
audio = empty_audio
|
||||
|
||||
# Stack frames to create video tensor [B, H, W, C]
|
||||
video_tensor = torch.stack(frames)[None, ...]
|
||||
|
||||
# Clean up temporary files
|
||||
try:
|
||||
os.unlink(video_path)
|
||||
for frame_file in frame_files:
|
||||
os.unlink(os.path.join(frames_dir, frame_file))
|
||||
os.rmdir(frames_dir)
|
||||
except Exception as e:
|
||||
print(f"Error cleaning up temp files: {e}")
|
||||
|
||||
print(f"Processed video with {len(frames)} frames")
|
||||
return (video_tensor, audio)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing video: {e}")
|
||||
# Fall back to empty frame and audio if video processing fails
|
||||
fallback_tensor = torch.zeros((1, 480, 640, 3))
|
||||
return (fallback_tensor, empty_audio)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unrecognized data URL format: {image[:30]}...")
|
||||
else:
|
||||
raise ValueError("Invalid image format received from webcam")
|
||||
raise ValueError("Invalid data received from webcam")
|
||||
|
||||
|
||||
coco_classes = [
|
||||
|
||||
+87
-17
@@ -9,6 +9,8 @@ app.registerExtension({
|
||||
// Store webcam stream at extension level to avoid reinitializing
|
||||
let webcamStream = null;
|
||||
let webcamVideo = null;
|
||||
let mediaRecorder = null;
|
||||
let recordedChunks = [];
|
||||
|
||||
// Initialize webcam once for all nodes
|
||||
const initWebcam = async () => {
|
||||
@@ -17,7 +19,7 @@ app.registerExtension({
|
||||
try {
|
||||
webcamStream = await navigator.mediaDevices.getUserMedia({
|
||||
video: true,
|
||||
audio: false
|
||||
audio: true // Now requesting audio too for video recording
|
||||
});
|
||||
// console.log("Got media stream:", webcamStream);
|
||||
|
||||
@@ -40,6 +42,60 @@ app.registerExtension({
|
||||
}
|
||||
return { stream: webcamStream, video: webcamVideo };
|
||||
};
|
||||
|
||||
// Function to record video for a specified duration
|
||||
const recordVideo = async (duration) => {
|
||||
return new Promise(async (resolve, reject) => {
|
||||
try {
|
||||
// Initialize webcam if not already done
|
||||
const { stream, video } = await initWebcam();
|
||||
|
||||
// Clear previous recording data
|
||||
recordedChunks = [];
|
||||
|
||||
// Initialize MediaRecorder
|
||||
mediaRecorder = new MediaRecorder(stream, {
|
||||
mimeType: 'video/webm;codecs=vp9,opus'
|
||||
});
|
||||
|
||||
// Handle data available event
|
||||
mediaRecorder.ondataavailable = (event) => {
|
||||
if (event.data.size > 0) {
|
||||
recordedChunks.push(event.data);
|
||||
}
|
||||
};
|
||||
|
||||
// Handle recording stop event
|
||||
mediaRecorder.onstop = () => {
|
||||
// Create a blob from recorded chunks
|
||||
const blob = new Blob(recordedChunks, { type: 'video/webm' });
|
||||
|
||||
// Convert blob to base64 data URL
|
||||
const reader = new FileReader();
|
||||
reader.onloadend = () => {
|
||||
const base64data = reader.result;
|
||||
resolve(base64data);
|
||||
};
|
||||
reader.readAsDataURL(blob);
|
||||
};
|
||||
|
||||
// Start recording
|
||||
mediaRecorder.start();
|
||||
console.log(`Recording started for ${duration} seconds`);
|
||||
|
||||
// Stop recording after specified duration
|
||||
setTimeout(() => {
|
||||
if (mediaRecorder && mediaRecorder.state !== 'inactive') {
|
||||
mediaRecorder.stop();
|
||||
console.log('Recording stopped');
|
||||
}
|
||||
}, duration * 1000);
|
||||
} catch (error) {
|
||||
console.error("Error during video recording:", error);
|
||||
reject(error);
|
||||
}
|
||||
});
|
||||
};
|
||||
|
||||
nodeType.prototype.getCustomWidgets = function() {
|
||||
// console.log("Getting custom widgets");
|
||||
@@ -69,8 +125,6 @@ app.registerExtension({
|
||||
const canvas = document.createElement("canvas");
|
||||
|
||||
camera.serializeValue = async () => {
|
||||
// console.log("Serializing camera value");
|
||||
|
||||
// Only initialize webcam when actually needed
|
||||
if (!webcamVideo || webcamVideo.paused || webcamVideo.readyState !== 4) {
|
||||
// console.log("Initializing webcam for first use...");
|
||||
@@ -78,21 +132,37 @@ app.registerExtension({
|
||||
webcamVideo = video;
|
||||
}
|
||||
|
||||
console.log("Capturing frame from video:", {
|
||||
videoWidth: webcamVideo.videoWidth,
|
||||
videoHeight: webcamVideo.videoHeight
|
||||
});
|
||||
// Find record_seconds input value
|
||||
const recordSecondsWidget = this.widgets.find(w => w.name === "record_seconds");
|
||||
const recordSeconds = recordSecondsWidget ? recordSecondsWidget.value : 0;
|
||||
|
||||
// Always use native webcam resolution
|
||||
canvas.width = webcamVideo.videoWidth;
|
||||
canvas.height = webcamVideo.videoHeight;
|
||||
|
||||
const ctx = canvas.getContext("2d");
|
||||
ctx.drawImage(webcamVideo, 0, 0, canvas.width, canvas.height);
|
||||
|
||||
const dataUrl = canvas.toDataURL("image/png");
|
||||
console.log("Captured frame at native resolution:", canvas.width, "x", canvas.height);
|
||||
return dataUrl;
|
||||
// Single image capture mode
|
||||
if (recordSeconds <= 0) {
|
||||
console.log("Capturing single frame from video:", {
|
||||
videoWidth: webcamVideo.videoWidth,
|
||||
videoHeight: webcamVideo.videoHeight
|
||||
});
|
||||
|
||||
// Always use native webcam resolution
|
||||
canvas.width = webcamVideo.videoWidth;
|
||||
canvas.height = webcamVideo.videoHeight;
|
||||
|
||||
const ctx = canvas.getContext("2d");
|
||||
ctx.drawImage(webcamVideo, 0, 0, canvas.width, canvas.height);
|
||||
|
||||
const dataUrl = canvas.toDataURL("image/png");
|
||||
console.log("Captured frame at native resolution:", canvas.width, "x", canvas.height);
|
||||
return dataUrl;
|
||||
}
|
||||
// Video recording mode
|
||||
else {
|
||||
console.log(`Starting video recording for ${recordSeconds} seconds`);
|
||||
|
||||
// Record video and return data URL
|
||||
const videoDataUrl = await recordVideo(recordSeconds);
|
||||
console.log("Video recording completed, sending data to backend");
|
||||
return videoDataUrl;
|
||||
}
|
||||
};
|
||||
|
||||
return r;
|
||||
|
||||
Reference in New Issue
Block a user