Created
June 23, 2025 11:28
-
-
Save yemeen/1871127992a87db6bba855128a114738 to your computer and use it in GitHub Desktop.
radial ect
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from matplotlib.colors import Normalize | |
| from pathlib import Path | |
| from ect import ECT, EmbeddedGraph | |
| from ect.results import ECTResult | |
| from data_processing.data_loader import create_leaf_dataframe, DatasetVariant | |
| import multiprocessing as mp | |
| from functools import partial | |
| def ect_to_radial(ect_matrix: ECTResult, title=None, save_path=None, show=False): | |
| """ | |
| Convert an ECT matrix to a radial plot. | |
| Args: | |
| ect_matrix: ECTResult object containing the matrix and metadata | |
| title: optional string for plot title | |
| save_path: optional path to save the figure | |
| show: whether to display the plot | |
| """ | |
| thetas = ect_matrix.directions.thetas | |
| thresholds = ect_matrix.thresholds | |
| THETA, R = np.meshgrid(thetas, thresholds) | |
| fig, ax = plt.subplots(subplot_kw=dict(projection="polar"), figsize=(10, 10)) | |
| im = ax.pcolormesh(THETA, R, ect_matrix.T, cmap="viridis") | |
| ax.set_theta_zero_location("N") | |
| ax.set_theta_direction(-1) | |
| plt.colorbar(im, label="ECT Value") | |
| if title: | |
| plt.title(title) | |
| plt.tight_layout() | |
| if save_path: | |
| save_path = Path(save_path) | |
| save_path.parent.mkdir(exist_ok=False, parents=True) | |
| plt.savefig(save_path, bbox_inches="tight", dpi=300) | |
| if show: | |
| plt.show() | |
| else: | |
| plt.close() | |
| def plot_leaf_and_radial( | |
| points, ect_result: ECTResult, title=None, save_path=None, show=False | |
| ): | |
| """ | |
| Create a figure with leaf outline, regular ECT, and radial ECT plots side by side. | |
| Args: | |
| points: numpy array of shape (N, 2) containing ALREADY SCALED leaf boundary points | |
| ect_result: ECTResult object containing the ECT data | |
| title: optional string for plot title | |
| save_path: optional path to save the figure | |
| show: whether to display the plot | |
| """ | |
| fig = plt.figure(figsize=(25, 8)) | |
| ax_leaf = plt.subplot(131) | |
| ax_leaf.plot(points[:, 0], points[:, 1], "k-", linewidth=2) | |
| ax_leaf.set_title("Leaf Outline") | |
| ax_leaf.axis("equal") | |
| ax_leaf.grid(True, alpha=0.3) | |
| ax_ect = plt.subplot(132) | |
| im_ect = ax_ect.imshow( | |
| ect_result.T, | |
| origin="lower", | |
| aspect="auto", | |
| extent=[0, 360, ect_result.thresholds[0], ect_result.thresholds[-1]], | |
| cmap="viridis", | |
| ) | |
| ax_ect.set_title("Regular ECT") | |
| ax_ect.set_xlabel("Direction (degrees)") | |
| ax_ect.set_ylabel("Threshold") | |
| plt.colorbar(im_ect, ax=ax_ect, label="ECT Value") | |
| ax_radial = plt.subplot(133, projection="polar") | |
| THETA, R = np.meshgrid(ect_result.directions.thetas, ect_result.thresholds) | |
| im_radial = ax_radial.pcolormesh(THETA, R, ect_result.T, cmap="viridis", alpha=1) | |
| x = points[:, 0] | |
| y = points[:, 1] | |
| r = np.sqrt(x**2 + y**2) | |
| theta = np.arctan2(y, x) | |
| max_r = np.max(ect_result.thresholds) | |
| scaled_r = r * max_r | |
| ax_radial.plot(theta, scaled_r, "-", color="white", linewidth=2, alpha=0.5) | |
| ax_radial.set_theta_zero_location("N") | |
| ax_radial.set_theta_direction(-1) | |
| ax_radial.set_title("Radial ECT with Leaf Overlay") | |
| plt.colorbar(im_radial, ax=ax_radial, label="ECT Value") | |
| if title: | |
| plt.suptitle(title, y=1.05) | |
| plt.tight_layout() | |
| if save_path: | |
| save_path = Path(save_path) | |
| save_path.parent.mkdir(exist_ok=True, parents=True) | |
| plt.savefig(save_path, bbox_inches="tight", dpi=300) | |
| if show: | |
| plt.show() | |
| else: | |
| plt.close() | |
| def process_sample(sample, class_name, class_dir, ect, idx): | |
| """Process a single leaf sample""" | |
| points = np.load(sample["file"]) | |
| G = EmbeddedGraph() | |
| G.add_cycle(points) | |
| G.center_coordinates(center_type="origin") | |
| G.transform_coordinates() | |
| G.scale_coordinates(1) | |
| ect_result = ect.calculate(G) | |
| scaled_points = G.coord_matrix | |
| leaf_name = Path(sample["file"]).stem | |
| save_path = class_dir / f"{leaf_name}_radial.png" | |
| plot_leaf_and_radial( | |
| scaled_points, | |
| ect_result, | |
| title=f"Leaf Shape and ECT Visualization\nClass: {class_name}, Sample {idx + 1}", | |
| save_path=save_path, | |
| show=False, | |
| ) | |
| return f"Processed {class_name} sample {idx + 1}/10" | |
| def main(): | |
| num_dirs = 90 | |
| thresholds = np.linspace(0, 1, num_dirs) | |
| ect = ECT(num_dirs=num_dirs, thresholds=thresholds, bound_radius=1) | |
| base_dir = Path("analysis/radial_ect") | |
| if base_dir.exists(): | |
| import shutil | |
| shutil.rmtree(base_dir) | |
| base_dir.mkdir(parents=True, exist_ok=False) | |
| dataset = create_leaf_dataframe(360, 360, data_variant=DatasetVariant.ORIGINAL) | |
| # create a pool of workers to speed things up | |
| num_cores = mp.cpu_count() | |
| pool = mp.Pool(num_cores) | |
| for class_name in sorted(dataset["group"].unique()): | |
| class_dir = base_dir / class_name | |
| class_dir.mkdir(exist_ok=True) | |
| class_samples = dataset[dataset["group"] == class_name].sample( | |
| n=10, random_state=42 | |
| ) | |
| results = pool.starmap( | |
| process_sample, | |
| [ | |
| (row[1], class_name, class_dir, ect, idx) | |
| for idx, row in enumerate(class_samples.iterrows()) | |
| ], | |
| ) | |
| for result in results: | |
| print(result) | |
| print(f"Completed class {class_name}") | |
| pool.close() | |
| pool.join() | |
| if __name__ == "__main__": | |
| main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment