From 71e0ee69d7314a99595d5c88e883d9116a721691 Mon Sep 17 00:00:00 2001 From: Sean Lynch Date: Sat, 2 Sep 2023 20:28:09 -0400 Subject: [PATCH] Initial commit --- LICENSE | 2 +- README.md | 34 +++++++ __init__.py | 2 + optical_flow.py | 260 ++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 297 insertions(+), 1 deletion(-) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 optical_flow.py diff --git a/LICENSE b/LICENSE index b5a6a28..e03a020 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2023 Sean Lynch +Copyright (c) 2022 the comfyui-optical-flow authors Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md new file mode 100644 index 0000000..2fdee19 --- /dev/null +++ b/README.md @@ -0,0 +1,34 @@ += About = + +This package contains three nodes to help you compute optical flow +between pairs of images, usually adjacent frames in a video, visualize +the flow, and apply the flow to another image of the same dimensions. + +Most of the code is from [Deforum](https://deforum.github.io/), so +this is released under the same license (MIT). + += Nodes = + +== Compute optical flow == + +This node takes two images, prev and current, and computes the optical +flow between them using either the DIS (Dense Inverse Search) medium +or fine method, or Farneback. The images must have the same +dimensions. + +== Apply optical flow == + +This node takes an image and applies an optical flow to it, so that +the motion matches the original image. This can be used for example to +improve consistency between video frames in a vid2vid workflow, by +applying the motion between the previous input frame and the current +one to the previous output frame before using it as input to a sampler. + +== Visualize optical flow == + +This node takes an image and a flow and produces an image visualizing +the flow on top of the image. The image must be the same size as the +images used to compute the flow in the first place. It's up to you +whether you use the "prev", "current", or an image you intend to apply +the flow to. + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..75f57b6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,2 @@ +from .optical_flow import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/optical_flow.py b/optical_flow.py new file mode 100644 index 0000000..13078f8 --- /dev/null +++ b/optical_flow.py @@ -0,0 +1,260 @@ +import numpy as np +import cv2 +import torch + + +def np2tensor(img_np: np.ndarray | list[np.ndarray]) -> torch.Tensor: + if isinstance(img_np, list): + return torch.cat([np2tensor(img) for img in img_np], dim=0) + + return torch.from_numpy(img_np.astype(np.float32) / 255.0).unsqueeze(0) + + +def tensor2np(tensor: torch.Tensor) -> list[np.ndarray]: + batch_count = tensor.size(0) if len(tensor.shape) > 3 else 1 + if batch_count > 1: + out = [] + for i in range(batch_count): + out.extend(tensor2np(tensor[i])) + return out + + return [np.clip(255.0 * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)] + + +def remap(img, flow, border_mode = cv2.BORDER_REFLECT_101): + # copyMakeBorder doesn't support wrap, but supports replicate. Replaces wrap with reflect101. + if border_mode == cv2.BORDER_WRAP: + border_mode = cv2.BORDER_REFLECT_101 + h, w = img.shape[:2] + displacement = int(h * 0.25), int(w * 0.25) + larger_img = cv2.copyMakeBorder(img, displacement[0], displacement[0], displacement[1], displacement[1], border_mode) + lh, lw = larger_img.shape[:2] + larger_flow = extend_flow(flow, lw, lh) + remapped_img = cv2.remap(larger_img, larger_flow, None, cv2.INTER_LINEAR, border_mode) + output_img = center_crop_image(remapped_img, w, h) + return output_img + + +def center_crop_image(img, w, h): + y, x, _ = img.shape + width_indent = int((x - w) / 2) + height_indent = int((y - h) / 2) + cropped_img = img[height_indent:y-height_indent, width_indent:x-width_indent] + return cropped_img + + +def extend_flow(flow, w, h): + # Get the shape of the original flow image + flow_h, flow_w = flow.shape[:2] + # Calculate the position of the image in the new image + x_offset = int((w - flow_w) / 2) + y_offset = int((h - flow_h) / 2) + # Generate the X and Y grids + x_grid, y_grid = np.meshgrid(np.arange(w), np.arange(h)) + # Create the new flow image and set it to the X and Y grids + new_flow = np.dstack((x_grid, y_grid)).astype(np.float32) + # Shift the values of the original flow by the size of the border + flow[:,:,0] += x_offset + flow[:,:,1] += y_offset + # Overwrite the middle of the grid with the original flow + new_flow[y_offset:y_offset+flow_h, x_offset:x_offset+flow_w, :] = flow + # Return the extended image + return new_flow + + +def get_flow_from_images(i1, i2, method, prev_flow=None): + if method == "DIS Medium": + flow = get_flow_from_images_DIS(i1, i2, 'medium', prev_flow) + elif method == "DIS Fine": + flow = get_flow_from_images_DIS(i1, i2, 'fine', prev_flow) + elif method == "Farneback": # Farneback Normal: + flow = get_flow_from_images_Farneback(i1, i2, prev_flow) + else: + # if we reached this point, something went wrong. raise an error: + raise RuntimeError(f"Invald flow method name: '{method}'") + + return flow + + +def get_flow_from_images_DIS(i1, i2, preset, prev_flow): + # DIS PRESETS CHART KEY: finest scale, grad desc its, patch size + # DIS_MEDIUM: 1, 25, 8 | DIS_FAST: 2, 16, 8 | DIS_ULTRAFAST: 2, 12, 8 + if preset == 'medium': preset_code = cv2.DISOPTICAL_FLOW_PRESET_MEDIUM + elif preset == 'fast': preset_code = cv2.DISOPTICAL_FLOW_PRESET_FAST + elif preset == 'ultrafast': preset_code = cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST + elif preset in ['slow','fine']: preset_code = None + i1 = cv2.cvtColor(i1, cv2.COLOR_BGR2GRAY) + i2 = cv2.cvtColor(i2, cv2.COLOR_BGR2GRAY) + dis = cv2.DISOpticalFlow_create(preset_code) + # custom presets + if preset == 'slow': + dis.setGradientDescentIterations(192) + dis.setFinestScale(1) + dis.setPatchSize(8) + dis.setPatchStride(4) + if preset == 'fine': + dis.setGradientDescentIterations(192) + dis.setFinestScale(0) + dis.setPatchSize(8) + dis.setPatchStride(4) + return dis.calc(i1, i2, prev_flow) + + +def get_flow_from_images_Farneback(i1, i2, preset="normal", last_flow=None, pyr_scale = 0.5, levels = 3, winsize = 15, iterations = 3, poly_n = 5, poly_sigma = 1.2, flags = 0): + flags = cv2.OPTFLOW_FARNEBACK_GAUSSIAN # Specify the operation flags + pyr_scale = 0.5 # The image scale (<1) to build pyramids for each image + if preset == "fine": + levels = 13 # The number of pyramid layers, including the initial image + winsize = 77 # The averaging window size + iterations = 13 # The number of iterations at each pyramid level + poly_n = 15 # The size of the pixel neighborhood used to find polynomial expansion in each pixel + poly_sigma = 0.8 # The standard deviation of the Gaussian used to smooth derivatives used as a basis for the polynomial expansion + else: # "normal" + levels = 5 # The number of pyramid layers, including the initial image + winsize = 21 # The averaging window size + iterations = 5 # The number of iterations at each pyramid level + poly_n = 7 # The size of the pixel neighborhood used to find polynomial expansion in each pixel + poly_sigma = 1.2 # The standard deviation of the Gaussian used to smooth derivatives used as a basis for the polynomial expansion + i1 = cv2.cvtColor(i1, cv2.COLOR_BGR2GRAY) + i2 = cv2.cvtColor(i2, cv2.COLOR_BGR2GRAY) + flags = 0 # flags = cv2.OPTFLOW_USE_INITIAL_FLOW + flow = cv2.calcOpticalFlowFarneback(i1, i2, last_flow, pyr_scale, levels, winsize, iterations, poly_n, poly_sigma, flags) + return flow + + +def image_transform_optical_flow(img, flow, border_mode=cv2.BORDER_REPLICATE, flow_reverse=False): + if not flow_reverse: + flow = -flow + h, w = img.shape[:2] + flow[:, :, 0] += np.arange(w) + flow[:, :, 1] += np.arange(h)[:,np.newaxis] + return remap(img, flow, border_mode) + + +def draw_flow_lines_in_grid_in_color(img, flow, step=8, magnitude_multiplier=1, min_magnitude = 0, max_magnitude = 10000): + flow = flow * magnitude_multiplier + h, w = img.shape[:2] + y, x = np.mgrid[step/2:h:step, step/2:w:step].reshape(2,-1).astype(int) + fx, fy = flow[y,x].T + lines = np.vstack([x, y, x+fx, y+fy]).T.reshape(-1, 2, 2) + lines = np.int32(lines + 0.5) + vis = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + vis = cv2.cvtColor(vis, cv2.COLOR_GRAY2BGR) + + mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1]) + hsv = np.zeros((flow.shape[0], flow.shape[1], 3), dtype=np.uint8) + hsv[...,0] = ang*180/np.pi/2 + hsv[...,1] = 255 + hsv[...,2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX) + bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) + vis = cv2.add(vis, bgr) + + # Iterate through the lines + for (x1, y1), (x2, y2) in lines: + # Calculate the magnitude of the line + magnitude = np.sqrt((x2 - x1)**2 + (y2 - y1)**2) + + # Only draw the line if it falls within the magnitude range + if min_magnitude <= magnitude <= max_magnitude: + b = int(bgr[y1, x1, 0]) + g = int(bgr[y1, x1, 1]) + r = int(bgr[y1, x1, 2]) + color = (b, g, r) + cv2.arrowedLine(vis, (x1, y1), (x2, y2), color, thickness=1, tipLength=0.1) + return vis + + +def visualize_flow(flow_img, flow): + flow_img = cv2.cvtColor(flow_img, cv2.COLOR_RGB2GRAY) + flow_img = cv2.cvtColor(flow_img, cv2.COLOR_GRAY2BGR) + flow_img = draw_flow_lines_in_grid_in_color(flow_img, flow) + return cv2.cvtColor(flow_img, cv2.COLOR_BGR2RGB) + + +class ComputeOpticalFlow: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prev": ("IMAGE",), + "current": ("IMAGE",), + "method": ([ + "DIS Medium", + "DIS Fine", + "Farneback", + ],), + }, + } + + RETURN_TYPES = ("OPTICAL_FLOW",) + FUNCTION = "compute_flow" + CATEGORY = "Optical flow" + + def compute_flow(self, prev, current, method): + images = zip(tensor2np(prev), tensor2np(current)) + return ([get_flow_from_images(im1, im2, method) for im1, im2 in images],) + + +class ApplyOpticalFlow: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "flow": ("OPTICAL_FLOW",), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "apply_flow" + CATEGORY = "Optical flow" + + def apply_flow(self, image, flow): + ims = tensor2np(image) + out = [image_transform_optical_flow(im, f) for im, f in zip(ims, flow)] + return (np2tensor(out),) + + +class VisualizeOpticalFlow: + """Visualize a flow as a set of arrows superimposed on the original image.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "flow": ("OPTICAL_FLOW",), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "visualize_flow" + CATEGORY = "Optical flow" + + def visualize_flow(self, image, flow): + ifs = zip(tensor2np(image), flow) + out = [visualize_flow(img, flow) for img, flow in ifs] + return (np2tensor(out),) + + +# A dictionary that contains all nodes you want to export with their names +# NOTE: names should be globally unique +NODE_CLASS_MAPPINGS = { + "Compute optical flow": ComputeOpticalFlow, + "Apply optical flow": ApplyOpticalFlow, + "Visualize optical flow": VisualizeOpticalFlow, +} + + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "ComputeOpticalFlow": "Compute optical flow", + "ApplyOpticalFlow": "Apply optical flow", + "VisualizeOpticalFlow": "Visualize optical flow", +}