diff --git a/nodes/core.py b/nodes/core.py index 11a1ca5..434d0cc 100644 --- a/nodes/core.py +++ b/nodes/core.py @@ -9,6 +9,8 @@ import numpy as np import io from PIL import Image +import torchvision.transforms as TT + logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') @@ -224,11 +226,12 @@ class KfCurveDraw: def INPUT_TYPES(cls): return { "required": { - "curve": ("KEYFRAMED_CURVE",) + "curve": ("KEYFRAMED_CURVE", {"forceInput": True,}), + "n": ("INT", {"default": 64}), } } - def main(self, curve): + def main(self, curve, n): """ """ @@ -238,25 +241,79 @@ class KfCurveDraw: # Build the plot using the provided function #build_plot(ax) #curve.plot(ax=ax) - curve.plot() - width, height = 10, 5 #inches + #curve.plot(n=n) + + eps:float=1e-9 + # value to be subtracted from keyframe to produce additional points for plotting. + # Plotting these additional values is important for e.g. visualizing step function behavior. + + m=3 + if n < m: + n = self.duration + 1 + n = max(m, n) + + + xs_base = list(range(int(n))) + list(curve.keyframes) + logger.debug(f"xs_base:{xs_base}") + xs = set() + for x in xs_base: + xs.add(x) + xs.add(x-eps) + + xs = [x for x in list(set(xs)) if (x >= 0)] + xs.sort() + ys = [curve[x] for x in xs] + + width, height = 12,8 #inches plt.figure(figsize=(width, height)) + #line = plt.plot(xs, ys, *args, **kargs) + line = plt.plot(xs, ys) + kfx = curve.keyframes + kfy = [curve[x] for x in kfx] + plt.scatter(kfx, kfy, color=line[0].get_color()) + + + + + #width, height = 10, 5 #inches + #plt.figure(figsize=(width, height)) # Save the plot to a BytesIO object buf = io.BytesIO() plt.savefig(buf, format='png', bbox_inches='tight') + plt.close() # no idea if this makes a difference buf.seek(0) # Read the image into a numpy array, converting it to RGB mode pil_image = Image.open(buf).convert('RGB') - plot_array = np.array(pil_image) #.astype(np.uint8) + #plot_array = np.array(pil_image) #.astype(np.uint8) # Convert the array to the desired shape [batch, channels, width, height] #plot_array = np.transpose(plot_array, (2, 0, 1)) # Reorder to [channels, width, height] #plot_array = np.expand_dims(plot_array, axis=0) # Add the batch dimension #plot_array = torch.tensor(plot_array) #.float() - plot_array = torch.from_numpy(plot_array) - return (plot_array,) + #plot_array = torch.from_numpy(plot_array) + + img_tensor = TT.ToTensor()(pil_image) + img_tensor = img_tensor.unsqueeze(0) + img_tensor = img_tensor.permute([0, 2, 3, 1]) + return (img_tensor,) + #return (plot_array,) + +# buffer_io = BytesIO() +# plt.savefig(buffer_io, format='png', bbox_inches='tight') +# plt.close() + +# buffer_io.seek(0) +# img = Image.open(buffer_io) + +# img_tensor = TT.ToTensor()(img) + +# img_tensor = img_tensor.unsqueeze(0) + +# img_tensor = img_tensor.permute([0, 2, 3, 1]) + +# return (img_tensor,) ###########################################