diff --git a/src/nodes/nodes_img.py b/src/nodes/nodes_img.py index 4497d27..20d6ece 100644 --- a/src/nodes/nodes_img.py +++ b/src/nodes/nodes_img.py @@ -1516,25 +1516,12 @@ class PlotMetricCurvesPIL(ComfyNodeABC): return img - def execute(self, metrics, plot_title, curve_color, width, height, auto_scale, legend_labels=None): - # --- Aggregate Data (same as matplotlib version) --- - plot_title_str = plot_title[0] - color_str = curve_color[0] - auto_scale = auto_scale[0] - - if not metrics: + def plot_one_dataset(self, datasets, plot_title_str, color_str, width, height, auto_scale, legend_labels): + if not datasets: logger.warning("No metrics data provided. Returning blank images.") - blank_image = pil_to_tensor_B(Image.new('RGB', (width[0], height[0]), 'white')) + blank_image = pil_to_tensor(Image.new('RGB', (width, height), 'white')) return (blank_image, blank_image) - # --- Normalize Input to List of Datasets --- - # Check if the input is a single dataset (List of Dicts) or multiple (List of Lists of Dicts) - # We assume if the first item is a dict, it's a single dataset. - if isinstance(metrics[0], dict): - datasets = [metrics] - else: - datasets = metrics - # --- Prepare Colors --- colors = color_str.split() n_colors = len(colors) @@ -1546,18 +1533,18 @@ class PlotMetricCurvesPIL(ComfyNodeABC): colors += [colors[-1]] * (n_datasets - n_colors) # --- Prepare Labels --- - labels = None - if legend_labels is not None: - # Ensure it's a list and has enough entries - # legend_labels comes in as a list of strings usually if from a primitive - if isinstance(legend_labels, list): - labels = legend_labels - else: - labels = [legend_labels] # Handle single string edge case - - # Pad labels if fewer than datasets - if len(labels) < n_datasets: - labels += [f"Curve {i+1}" for i in range(len(labels), n_datasets)] + labels = legend_labels # None +# if legend_labels is not None: +# # Ensure it's a list and has enough entries +# # legend_labels comes in as a list of strings usually if from a primitive +# if isinstance(legend_labels, list): +# labels = legend_labels +# else: +# labels = [legend_labels] # Handle single string edge case +# +# # Pad labels if fewer than datasets +# if len(labels) < n_datasets: +# labels += [f"Curve {i+1}" for i in range(len(labels), n_datasets)] # --- Process Each Dataset --- pr_curves_data = [] # Will hold tuples: (x_data, y_data, color) @@ -1601,7 +1588,7 @@ class PlotMetricCurvesPIL(ComfyNodeABC): if not pr_curves_data: logger.warning("Metrics data did not contain F/FP/FR keys. Returning blank images.") - blank_image = pil_to_tensor_B(Image.new('RGB', (width[0], height[0]), 'white')) + blank_image = pil_to_tensor(Image.new('RGB', (width, height), 'white')) return (blank_image, blank_image) # Determine limits @@ -1618,22 +1605,87 @@ class PlotMetricCurvesPIL(ComfyNodeABC): pr_plot_pil = self._create_plot_with_pil( curves_data=pr_curves_data, title=f"{plot_title_str} (PR Curve)", x_label="Recall", y_label="Precision", - width=width[0], height=height[0], + width=width, height=height, x_lim=pr_x_lim, y_lim=pr_y_lim ) - pr_plot_tensor = pil_to_tensor_B(pr_plot_pil) + pr_plot_tensor = pil_to_tensor(pr_plot_pil) # --- Generate F-Measure Curve Plot --- # F-measure is always 0-255 on X and 0-1 on Y fm_plot_pil = self._create_plot_with_pil( curves_data=fm_curves_data, title=f"{plot_title_str} (F-Measure Curve)", x_label="Threshold", y_label="F-measure", - width=width[0], height=height[0], + width=width, height=height, x_lim=(0, F_POINTS), y_lim=(0.0, 1.0) ) - fm_plot_tensor = pil_to_tensor_B(fm_plot_pil) + fm_plot_tensor = pil_to_tensor(fm_plot_pil) - return (pr_plot_tensor, fm_plot_tensor) + return pr_plot_tensor, fm_plot_tensor + + def normalize(self, param_name, param_values, n): + """ Ensure we have N values for a parameter """ + got = len(param_values) + if got > n: + logger.warning(f"Got {got} {param_name} and {n} titles, discarding extra values") + return param_values[:n] + if got == 1: + return param_values * n + logger.warning(f"Got {got} {param_name} and {n} titles, repeating the last one") + return param_values + param_values[-1:] * (n - got) + + def execute(self, metrics, plot_title, curve_color, width, height, auto_scale, legend_labels=None): + # --- Normalize Input to List of Datasets --- + # Check if the input is a single dataset (List of Dicts) or multiple (List of Lists of Dicts) + # We assume if the first item is a dict, it's a single dataset. + if isinstance(metrics[0], dict): + datasets = [metrics] + else: + datasets = metrics + n_datasets = len(datasets) + n_plot_title = len(plot_title) + + # We will plot one PR/FM pair for each title. + # We assume each title will have the same number of metrics. + if n_datasets < n_plot_title: + msg = f"Too much titles ({n_plot_title}) for {n_datasets} metrics" + logger.error(msg) + raise ValueError(msg) + if n_datasets % n_plot_title != 0: + msg = f"Got {n_plot_title} titles for {n_datasets} metrics, must be multiples" + logger.error(msg) + raise ValueError(msg) + slice_size = n_datasets // n_plot_title + + # Check the legend labels + if legend_labels is not None: + n_legend_labels = len(legend_labels) + if n_legend_labels < slice_size or (n_legend_labels > slice_size and n_legend_labels != n_datasets): + msg = f"Got {n_legend_labels} legend labels for {n_datasets} metrics" + logger.error(msg) + raise ValueError(msg) + if n_legend_labels != n_datasets: # Here we know n_legend_labels == slice_size + # Repeat the same for all titles + legend_labels = legend_labels * n_plot_title + + color = self.normalize("curve_color", curve_color, n_plot_title) + auto_scale = self.normalize("auto_scale", auto_scale, n_plot_title) + width = self.normalize("width", width, n_plot_title) + height = self.normalize("height", height, n_plot_title) + + logger.debug(f"Generating plot for {n_datasets} metrics and {n_plot_title} titles") + + prs = [] + fms = [] + for c, title in enumerate(plot_title): + logger.debug(f"- {title}") + start = c * slice_size + end = start + slice_size + pr, fm = self.plot_one_dataset(datasets[start:end], title, color[c], width[c], height[c], auto_scale[c], + legend_labels if legend_labels is None else legend_labels[start:end]) + prs.append(pr) + fms.append(fm) + + return (torch.stack(prs), torch.stack(fms)) class CompositeFace(ComfyNodeABC):