Merge pull request #26 from ryanontheinside/feat/fastwebcam-improvements
Feat/fastwebcam improvements
This commit is contained in:
@@ -49,52 +49,158 @@ class DTypeConverter:
|
||||
return (converted,)
|
||||
|
||||
|
||||
|
||||
class FastWebcamCapture:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("WEBCAM", {}),
|
||||
"width": ("INT", {"default": 640, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
||||
"height": ("INT", {"default": 480, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
||||
"capture_on_queue": ("BOOLEAN", {"default": True}),
|
||||
"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)."
|
||||
}
|
||||
),
|
||||
"capture_on_queue": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": True,
|
||||
"tooltip": "When enabled, capture new frames on each queue. When disabled, reuse previously captured video."
|
||||
}
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_TYPES = ("IMAGE", "INT")
|
||||
RETURN_NAMES = ("image", "frame_count")
|
||||
FUNCTION = "process_capture"
|
||||
|
||||
CATEGORY = "image"
|
||||
|
||||
def process_capture(self, image, width, height, capture_on_queue):
|
||||
# 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)
|
||||
def __init__(self):
|
||||
self.last_capture = None
|
||||
self.last_frame_count = 0
|
||||
|
||||
# Convert base64 to PIL Image
|
||||
buffer = BytesIO(base64.b64decode(base64_data))
|
||||
pil_image = Image.open(buffer).convert("RGB")
|
||||
def process_capture(self, image, record_seconds, capture_on_queue):
|
||||
"""
|
||||
Process a webcam image or video at 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.
|
||||
"""
|
||||
# If capture_on_queue is disabled and we have a previous capture, reuse it
|
||||
if not capture_on_queue and self.last_capture is not None:
|
||||
print("Reusing previous capture")
|
||||
return (self.last_capture, self.last_frame_count)
|
||||
|
||||
# Convert PIL to numpy array
|
||||
image = np.array(pil_image)
|
||||
import torch
|
||||
import numpy as np
|
||||
import json
|
||||
|
||||
# Check if we're receiving image or JSON frames data
|
||||
if isinstance(image, str):
|
||||
# Single image data URL
|
||||
if image.startswith("data:image/"):
|
||||
# Extract the base64 data after the comma
|
||||
base64_data = re.sub("^data:image/.+;base64,", "", image)
|
||||
|
||||
# Handle resize if requested
|
||||
if width > 0 and height > 0:
|
||||
import cv2
|
||||
# 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}")
|
||||
|
||||
image = cv2.resize(image, (width, height), interpolation=cv2.INTER_AREA)
|
||||
# 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, ...]
|
||||
|
||||
# Store for potential reuse
|
||||
self.last_capture = image_tensor
|
||||
self.last_frame_count = 1
|
||||
|
||||
return (image_tensor, 1)
|
||||
|
||||
# JSON frames data from sequence capture
|
||||
elif image.startswith("{") and "frames" in image:
|
||||
print("Received frame sequence data from webcam")
|
||||
|
||||
try:
|
||||
# Parse JSON data
|
||||
frame_data = json.loads(image)
|
||||
frames = frame_data.get("frames", [])
|
||||
fps = frame_data.get("fps", 30)
|
||||
duration = frame_data.get("duration", 1.0)
|
||||
|
||||
if not frames:
|
||||
raise ValueError("No frames received in frame data")
|
||||
|
||||
print(f"Processing {len(frames)} frames captured at {fps} FPS over {duration} seconds")
|
||||
|
||||
# Convert frames to tensor
|
||||
tensor_frames = []
|
||||
|
||||
for i, frame_dataurl in enumerate(frames):
|
||||
# Extract base64 data
|
||||
base64_data = re.sub("^data:image/.+;base64,", "", frame_dataurl)
|
||||
|
||||
# Convert to PIL image
|
||||
buffer = BytesIO(base64.b64decode(base64_data))
|
||||
pil_frame = Image.open(buffer).convert("RGB")
|
||||
|
||||
# Convert to numpy and normalize
|
||||
frame_np = np.array(pil_frame).astype(np.float32) / 255.0
|
||||
|
||||
# Add to frame list
|
||||
tensor_frames.append(torch.from_numpy(frame_np))
|
||||
|
||||
# Stack frames to create video tensor
|
||||
# ComfyUI expects image tensors in BHWC format (batch, height, width, channels)
|
||||
video_tensor = torch.stack(tensor_frames) # First get [frames, H, W, C]
|
||||
|
||||
print(f"Stacked tensor shape before corrections: {video_tensor.shape}")
|
||||
|
||||
# Then interpret the frames as the batch dimension for ComfyUI
|
||||
if len(video_tensor.shape) == 4:
|
||||
# This is already correct format [frames, H, W, C]
|
||||
pass
|
||||
elif len(video_tensor.shape) == 5:
|
||||
# Remove extra batch dimension if present
|
||||
video_tensor = video_tensor.squeeze(0)
|
||||
else:
|
||||
# Something's wrong with the shape
|
||||
raise ValueError(f"Unexpected tensor shape: {video_tensor.shape}. Expected 4D [frames, H, W, C]")
|
||||
|
||||
print(f"Final video tensor shape: {video_tensor.shape}")
|
||||
|
||||
# Store for potential reuse
|
||||
self.last_capture = video_tensor
|
||||
self.last_frame_count = len(frames)
|
||||
|
||||
return (video_tensor, len(frames))
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
print(f"Error processing frame sequence: {e}")
|
||||
print(traceback.format_exc())
|
||||
raise RuntimeError(f"Failed to process frame sequence: {str(e)}")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unrecognized data format from webcam")
|
||||
else:
|
||||
raise ValueError("Invalid image format received from webcam")
|
||||
raise ValueError("Invalid data received from webcam")
|
||||
|
||||
|
||||
coco_classes = [
|
||||
@@ -1318,3 +1424,4 @@ class RenormalizeInt:
|
||||
result = max(output_max, min(output_min, result))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
+114
-20
@@ -9,6 +9,7 @@ app.registerExtension({
|
||||
// Store webcam stream at extension level to avoid reinitializing
|
||||
let webcamStream = null;
|
||||
let webcamVideo = null;
|
||||
let lastCapture = null;
|
||||
|
||||
// Initialize webcam once for all nodes
|
||||
const initWebcam = async () => {
|
||||
@@ -17,7 +18,7 @@ app.registerExtension({
|
||||
try {
|
||||
webcamStream = await navigator.mediaDevices.getUserMedia({
|
||||
video: true,
|
||||
audio: false
|
||||
audio: false // No need for audio anymore
|
||||
});
|
||||
// console.log("Got media stream:", webcamStream);
|
||||
|
||||
@@ -40,6 +41,74 @@ app.registerExtension({
|
||||
}
|
||||
return { stream: webcamStream, video: webcamVideo };
|
||||
};
|
||||
|
||||
// Function to capture frames for a given duration
|
||||
const captureFrames = async (duration, fps = 30) => {
|
||||
return new Promise(async (resolve, reject) => {
|
||||
try {
|
||||
// Initialize webcam if not already done
|
||||
const { video } = await initWebcam();
|
||||
|
||||
// Create canvas for frame capture
|
||||
const canvas = document.createElement('canvas');
|
||||
canvas.width = video.videoWidth;
|
||||
canvas.height = video.videoHeight;
|
||||
const ctx = canvas.getContext('2d');
|
||||
|
||||
// Calculate total frames to capture
|
||||
const totalFrames = Math.ceil(duration * fps);
|
||||
const frameInterval = 1000 / fps; // ms between frames
|
||||
|
||||
console.log(`Capturing ${totalFrames} frames at ${fps} FPS (${frameInterval}ms intervals)`);
|
||||
|
||||
// Store captured frames
|
||||
const frames = [];
|
||||
|
||||
// Start capturing frames at regular intervals
|
||||
let frameCount = 0;
|
||||
const startTime = Date.now();
|
||||
|
||||
const captureFrame = () => {
|
||||
// Draw current video frame to canvas
|
||||
ctx.drawImage(video, 0, 0, canvas.width, canvas.height);
|
||||
const frameDataUrl = canvas.toDataURL('image/jpeg', 0.95);
|
||||
frames.push(frameDataUrl);
|
||||
|
||||
frameCount++;
|
||||
if (frameCount < totalFrames) {
|
||||
// Calculate next frame time based on start time to avoid drift
|
||||
const nextFrameTime = startTime + Math.round(frameCount * frameInterval);
|
||||
const delay = Math.max(0, nextFrameTime - Date.now());
|
||||
|
||||
// Schedule next frame capture
|
||||
setTimeout(captureFrame, delay);
|
||||
} else {
|
||||
// All frames captured
|
||||
console.log(`Completed capture of ${frames.length} frames`);
|
||||
|
||||
const result = {
|
||||
frames: frames,
|
||||
width: canvas.width,
|
||||
height: canvas.height,
|
||||
duration: duration,
|
||||
fps: fps
|
||||
};
|
||||
|
||||
// Store for potential reuse
|
||||
lastCapture = result;
|
||||
|
||||
resolve(result);
|
||||
}
|
||||
};
|
||||
|
||||
// Start capture process
|
||||
captureFrame();
|
||||
} catch (error) {
|
||||
console.error("Error during capture:", error);
|
||||
reject(error);
|
||||
}
|
||||
});
|
||||
};
|
||||
|
||||
nodeType.prototype.getCustomWidgets = function() {
|
||||
// console.log("Getting custom widgets");
|
||||
@@ -66,13 +135,15 @@ app.registerExtension({
|
||||
const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
|
||||
|
||||
const camera = this.widgets.find((w) => w.name === "image");
|
||||
const w = this.widgets.find((w) => w.name === "width");
|
||||
const h = this.widgets.find((w) => w.name === "height");
|
||||
|
||||
const canvas = document.createElement("canvas");
|
||||
|
||||
camera.serializeValue = async () => {
|
||||
// console.log("Serializing camera value");
|
||||
// Check if we should reuse previous capture
|
||||
const captureOnQueueWidget = this.widgets.find(w => w.name === "capture_on_queue");
|
||||
if (!captureOnQueueWidget.value && lastCapture) {
|
||||
console.log("Reusing previous capture");
|
||||
return JSON.stringify(lastCapture);
|
||||
}
|
||||
|
||||
// Only initialize webcam when actually needed
|
||||
if (!webcamVideo || webcamVideo.paused || webcamVideo.readyState !== 4) {
|
||||
@@ -81,22 +152,45 @@ app.registerExtension({
|
||||
webcamVideo = video;
|
||||
}
|
||||
|
||||
console.log("Capturing frame from video:", {
|
||||
paused: webcamVideo.paused,
|
||||
readyState: webcamVideo.readyState,
|
||||
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;
|
||||
|
||||
canvas.width = w.value || webcamVideo.videoWidth || 640;
|
||||
canvas.height = h.value || webcamVideo.videoHeight || 480;
|
||||
|
||||
const ctx = canvas.getContext("2d");
|
||||
ctx.drawImage(webcamVideo, 0, 0, canvas.width, canvas.height);
|
||||
|
||||
const dataUrl = canvas.toDataURL("image/png");
|
||||
// console.log("Successfully captured and converted frame to data URL");
|
||||
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");
|
||||
|
||||
// Store for potential reuse
|
||||
lastCapture = {
|
||||
frames: [dataUrl],
|
||||
width: canvas.width,
|
||||
height: canvas.height,
|
||||
duration: 0,
|
||||
fps: 30
|
||||
};
|
||||
|
||||
return dataUrl;
|
||||
}
|
||||
// Video capture mode - capture frames
|
||||
else {
|
||||
// Capture frames
|
||||
const captureResult = await captureFrames(recordSeconds);
|
||||
|
||||
// Convert to JSON and return
|
||||
return JSON.stringify(captureResult);
|
||||
}
|
||||
};
|
||||
|
||||
return r;
|
||||
|
||||
Reference in New Issue
Block a user