Nuke Copycat: Now Supports Auto Frame Selection for Sequences

import cv2
import numpy as np
import nuke
import OpenEXR
import Imath
import os

# Function to load frames from EXR files
def load_exr_frames(exr_directory):
    exr_files = sorted([f for f in os.listdir(exr_directory) if f.endswith('.exr')])
    frames = []
    for file_name in exr_files:
        file_path = os.path.join(exr_directory, file_name)
        exr_file = OpenEXR.InputFile(file_path)
        dw = exr_file.header()['dataWindow']
        size = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1)
        pt = Imath.PixelType(Imath.PixelType.FLOAT)
        rgb = [np.frombuffer(exr_file.channel(c, pt), dtype=np.float32) for c in ('R', 'G', 'B')]
        img_array = np.dstack(rgb).reshape((size[1], size[0], 3))
        rescaled_img_array = cv2.resize(img_array, (size[0] // 16, size[1] // 16), interpolation=cv2.INTER_AREA)
        normalized_img_array = cv2.normalize(rescaled_img_array, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        frames.append(normalized_img_array)
    return frames

# Determine the file type and load frames
n = nuke.selectedNode()
input_node = n.input(0)
first_frame = int(input_node['first'].value())  # Retrieve the first frame number
last_frame = int(input_node['last'].value())
if input_node and input_node.Class() == 'Read':
    file_path = input_node['file'].getValue()
	
    file_extension = os.path.splitext(file_path)[-1].lower()
    if file_extension in ['.mov', '.mp4']:
        cap = cv2.VideoCapture(file_path)
        frames = [cap.read()[1] for _ in range(int(input_node['last'].value()))]
        cap.release()
    elif file_extension == '.exr':
        exr_directory = os.path.dirname(file_path)
        frames = load_exr_frames(exr_directory)
    else:
        raise Exception("Unsupported file format")
else:
    raise Exception("No Read node connected or selected node is not a Read node")

if not frames:
    raise Exception("Failed to load frames from the file")

# Assume the maximum threshold value is 100
max_threshold_value = 100

# Get the user input from the knob
user_input_threshold = int(nuke.thisNode().knob('threshold').value())

# Calculate the inverse threshold
threshold = max_threshold_value - user_input_threshold


# Motion detection and frame processing

previous_frame = None

def detect_motion(previous_frame, current_frame, threshold):
    # Calculate the frame difference
    frame_diff = cv2.absdiff(previous_frame, current_frame)
    
    # Threshold the difference image to get a binary image
    _, thresholded_diff = cv2.threshold(frame_diff, threshold, 255, cv2.THRESH_BINARY)
    
    # Ensure the binary image is in 8-bit format, required for findContours
    thresholded_diff = thresholded_diff.astype(np.uint8)
    
    # Define a structuring element for morphological operations
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    
    # Perform morphological open operation to remove noise
    opened_diff = cv2.morphologyEx(thresholded_diff, cv2.MORPH_OPEN, kernel)
    
    # Find contours from the binary image
    contours, _ = cv2.findContours(opened_diff, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Calculate bounding rectangles for each contour
    return [cv2.boundingRect(contour) for contour in contours]

def process_frame(frame, frame_number):
    global previous_frame
    current_frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    if previous_frame is None:
        previous_frame = current_frame_gray
        return frame
    motion_regions = detect_motion(previous_frame, current_frame_gray, threshold)
    for x, y, w, h in motion_regions:
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
    cv2.putText(frame, f"Frame: {int(frame_number)}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
    previous_frame = current_frame_gray
    return frame

frame_sums = []
for frame_number, frame in enumerate(frames):
    if frame is None:
        print("Error reading frame")
        break
    processed_frame = process_frame(frame, frame_number)
    motion_regions = detect_motion(previous_frame, cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), threshold)
    frame_sum = sum(w * h for _, _, w, h in motion_regions)
    frame_sums.append(frame_sum)
    cv2.imshow('Processed Frame', processed_frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cv2.destroyAllWindows()

# Create FrameHold nodes based on motion
count = int(nuke.thisNode().knob('count').value())
read_node_x = n.xpos()
read_node_y = n.ypos()
indices = np.argsort(frame_sums)[-count:][::-1]  # Get indices of frames with the most motion

for i, frame_index in enumerate(indices):
    actual_frame_number = frame_index + first_frame  # Adjust index by adding first_frame
    frame_sum = frame_sums[frame_index]
    framehold_node = nuke.createNode('FrameHold')
    framehold_node['first_frame'].setValue(actual_frame_number)  # Set to actual frame number
    framehold_node['label'].setValue(f"Sum: {frame_sum}")
    node_x = read_node_x - ((count - 1) * 100 // 2) + 100 * i
    node_y = read_node_y + 150
    framehold_node.setXpos(node_x)
    framehold_node.setYpos(node_y)
    framehold_node.setInput(0, n)