diff --git a/README.md b/README.md
index 817155b..02af137 100644
--- a/README.md
+++ b/README.md
@@ -106,17 +106,26 @@ This allows for qualitative inspection of segmentation accuracy and typical erro
---
-### 6️⃣ Figure X
-IoU scores for all evaluated image sets.
-The chart shows segmentation performance for each image, measured as the overlap between the predicted mask and the ground truth.
-
-
- 
-
- Chart for comparison results(IoU)[img_comparison.png]
-
-
-
+### 6️⃣ Output of results in csv
+A table containing IoU scores for evaluated images is automatically saved.
+File:
+```
+iou_results.csv
+```
+Format:
+```
+image_name - iou
+RGB_ar037_2019_n_06_04_0.png - 52.10%
+RGB_ar037_2019_n_07_05_0.png - 17.80%
+RGB_ar037_2019_n_07_17_1.png - 0.00%
+RGB_ar037_2019_n_08_14_0.png - 41.98%
+RGB_ar037_2019_n_13_13_0.png - 15.57%
+RGB_ar039_2019_n_02_04_0.png - 50.59%
+RGB_ar039_2019_n_02_10_0.png - 13.10%
+RGB_ar039_2019_n_03_04_0.png - 14.69%
+RGB_ar039_2019_n_04_10_0.png - 8.10%
+RGB_ar039_2019_n_04_13_0.png - 44.99%
+```
---
---
@@ -177,14 +186,18 @@ python main.py --help
```bash
python main.py \
- -c --config config.yaml \
- -ni --num-images 5 \
- -nc --num-compare 5 \
- -h_min --hue-min 0.75 \
- -h_max --hue-max 0.95 \
- -s --sat-thr 0.25 \
- -v --val-thr 0.55 \
- -o --output-dir output/final_run/
+ -h, --help show this help message and exit
+ -c, --config CONFIG / Path to YAML configuration file (default: config.yaml)
+ -h_min, --hue_min HUE_MIN / Minimum HSV hue threshold for RGB segmentation
+ -h_max, --hue_max HUE_MAX / Maximum HSV hue threshold for RGB segmentation
+ -s, --sat_thr SAT_THR HSV / saturation threshold for RGB segmentation
+ -v, --val_thr VAL_THR HSV value / (brightness) threshold for RGB segmentation
+ -n_img, --num_images NUM_IMAGES / Number of images to load or preview from the dataset
+ -n_cmp, --num_compare NUM_COMPARE / Number of images used for quantitative evaluation
+ -o, --output_dir OUTPUT_DIR / Output directory for generated masks, metrics and logs
+ -p_rgb, --path_rgb PATH_RGB / Path or glob pattern to RGB images (e.g. data/RGB_images/*.png)
+ -p_nrg, --path_nrg PATH_NRG / Path or glob pattern to NRG images (e.g. data/NRG_images/*.png)
+ -p_mask, --path_mask PATH_MASK / Path or glob pattern to ground truth masks (e.g. data/masks/*.png)
```
> Command-line arguments have **higher priority** than values defined in `config.yaml`.
@@ -247,6 +260,11 @@ pip install -r requirements.txt
```bash
python main.py
```
+or with CLI overrides:
+
+```bash
+python main.py -n_cmp 5 -o data/generated
+```
Optional parameters can be supplied via the CLI.
@@ -255,19 +273,33 @@ Optional parameters can be supplied via the CLI.
## 📁 Project Structure
```
-Dead-Tree-Segmentation-main/
+Dead-Tree-Segmentation/
│
-├── main.py # Main pipeline script
-├── requirements.txt # Project dependencies
-├── temp_config.yaml # Configuration template
-├── config.yaml # Local config (gitignored)
+├── main.py
+├── requirements.txt
+├── temp_config.yaml
+├── config.yaml # gitignored
├── README.md
├── .gitignore
+│
+├── src/
+│ ├── cli.py
+│ ├── config_loader.py
+│ ├── segmentation.py
+│ ├── evaluation.py
+│ ├── io_utils.py
+│ └── logger.py
+│
├── data/
│ ├── RGB_images/
│ ├── NRG_images/
│ └── masks/
-└── output/ # Generated results
+│
+├── data_examples/
+│ ├── example_generated_mask/
+│ └── example_iou_results.csv
+│
+
```
---
diff --git a/config.yaml b/config.yaml
index 890984e..56e5b6a 100644
--- a/config.yaml
+++ b/config.yaml
@@ -2,7 +2,7 @@ paths:
rgb: ./data/RGB_images/*.png #
nrg: ./data/NRG_images/*.png #
masks: ./data/masks/*.png #
- output: ./data/generated_masks/ #
+ output: ./data_examples/example_generated_mask/ #
thresholds:
hue_min: 0.75
diff --git a/data/generated_mask/example_generated_mask/RGB_ar037_2019_n_06_04_0_mask.png b/data_examples/example_generated_mask/RGB_ar037_2019_n_06_04_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar037_2019_n_06_04_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar037_2019_n_06_04_0_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar037_2019_n_07_05_0_mask.png b/data_examples/example_generated_mask/RGB_ar037_2019_n_07_05_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar037_2019_n_07_05_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar037_2019_n_07_05_0_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar037_2019_n_07_17_1_mask.png b/data_examples/example_generated_mask/RGB_ar037_2019_n_07_17_1_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar037_2019_n_07_17_1_mask.png
rename to data_examples/example_generated_mask/RGB_ar037_2019_n_07_17_1_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar037_2019_n_08_14_0_mask.png b/data_examples/example_generated_mask/RGB_ar037_2019_n_08_14_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar037_2019_n_08_14_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar037_2019_n_08_14_0_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar037_2019_n_13_13_0_mask.png b/data_examples/example_generated_mask/RGB_ar037_2019_n_13_13_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar037_2019_n_13_13_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar037_2019_n_13_13_0_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar039_2019_n_02_04_0_mask.png b/data_examples/example_generated_mask/RGB_ar039_2019_n_02_04_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar039_2019_n_02_04_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar039_2019_n_02_04_0_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar039_2019_n_02_10_0_mask.png b/data_examples/example_generated_mask/RGB_ar039_2019_n_02_10_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar039_2019_n_02_10_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar039_2019_n_02_10_0_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar039_2019_n_03_04_0_mask.png b/data_examples/example_generated_mask/RGB_ar039_2019_n_03_04_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar039_2019_n_03_04_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar039_2019_n_03_04_0_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar039_2019_n_04_10_0_mask.png b/data_examples/example_generated_mask/RGB_ar039_2019_n_04_10_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar039_2019_n_04_10_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar039_2019_n_04_10_0_mask.png
diff --git a/data/generated_mask/example_generated_mask/RGB_ar039_2019_n_04_13_0_mask.png b/data_examples/example_generated_mask/RGB_ar039_2019_n_04_13_0_mask.png
similarity index 100%
rename from data/generated_mask/example_generated_mask/RGB_ar039_2019_n_04_13_0_mask.png
rename to data_examples/example_generated_mask/RGB_ar039_2019_n_04_13_0_mask.png
diff --git a/example_iou_results.csv b/data_examples/example_iou_results.csv
similarity index 100%
rename from example_iou_results.csv
rename to data_examples/example_iou_results.csv
diff --git a/main.py b/main.py
index a0f65cc..08fe2ad 100644
--- a/main.py
+++ b/main.py
@@ -1,572 +1,117 @@
-# Segmentation Pipeline
-
-#1. Load RGB and NRG images
-#2. Generate independent segmentation masks (RGB / NRG)
-#3. Fuse masks using morphological operations
-#4. Compare predicted masks with main masks
-#5. Compute IoU and confusion matrix
-#6. Optimize segmentation thresholds
-
-
-# Libraries
-
-import numpy as np
-import configparser
-from skimage import io, color, morphology, filters
-import matplotlib.pyplot as plt
-import os, glob, cv2, yaml, argparse
-from sklearn.metrics import confusion_matrix, classification_report
-
-# Loading configuration from config.yaml
-
-def load_config(path="config.yaml"):
- with open(path, "r", encoding="utf-8") as f:
- cfg = yaml.safe_load(f)
-
- config = {
- "paths_img": cfg["paths"]["rgb"],
- "paths_img_nrg": cfg["paths"]["nrg"],
- "paths_mask": cfg["paths"]["masks"],
- "output_dir": cfg["paths"]["output"],
-
- "hue_min": float(cfg["thresholds"]["hue_min"]),
- "hue_max": float(cfg["thresholds"]["hue_max"]),
- "sat_thr": float(cfg["thresholds"]["sat_thr"]),
- "val_thr": float(cfg["thresholds"]["val_thr"]),
-
- "num_images": int(cfg["general"]["num_images"]),
- "num_compare": int(cfg["general"]["num_compare"]),
- }
-
- return config
-
-def parse_args():
- parser = argparse.ArgumentParser(
- description="Dead Tree Segmentation Pipeline (RGB + NRG)"
- )
-
- parser.add_argument("-c", "--config", type=str, default="config.yaml",
- help="Path to config.yaml file")
-
- # Thresholds (override config)
- parser.add_argument("-h_min", "--hue-min", type=float, help="HSV hue min threshold")
- parser.add_argument("-h_max", "--hue-max", type=float, help="HSV hue max threshold")
- parser.add_argument("-s", "--sat-thr", type=float, help="HSV saturation threshold")
- parser.add_argument("-v", "--val-thr", type=float, help="HSV value threshold")
-
- # General
- parser.add_argument("-ni", "--num-images", type=int, help="Number of images to preview")
- parser.add_argument("-nc", "--num-compare", type=int, help="Number of images for evaluation")
-
- # Output
- parser.add_argument("-o", "--output-dir", type=str, help="Output directory for results")
-
- return parser.parse_args()
-
-
-args = parse_args()
-config = load_config(args.config)
-
-if args.hue_min is not None:
- config["hue_min"] = args.hue_min
-if args.hue_max is not None:
- config["hue_max"] = args.hue_max
-if args.sat_thr is not None:
- config["sat_thr"] = args.sat_thr
-if args.val_thr is not None:
- config["val_thr"] = args.val_thr
-
-if args.num_images is not None:
- config["num_images"] = args.num_images
-if args.num_compare is not None:
- config["num_compare"] = args.num_compare
+from src.cli import parse_args
+from src.config_loader import load_config
+from src.segmentation import (
+ generate_segment_mask_rgb,
+ generate_segment_mask_nrg,
+ fuse_masks
+)
+from src.evaluation import iou, compute_confusion_matrix
+from src.io_utils import save_iou_results, save_masks
-if args.output_dir is not None:
- config["output_dir"] = args.output_dir
+import glob, logging, cv2, os
+from skimage import io
-# Check if num_compare is greater than num_images
-if config["num_compare"] > config["num_images"]:
- raise ValueError(
- f"Invalid configuration: num_compare ({config['num_compare']}) "
- f"cannot be greater than num_images ({config['num_images']})."
+def main():
+ # === Logging setup ===
+ logging.basicConfig(
+ level=logging.INFO,
+ format="[%(levelname)s] %(message)s"
)
+ # === Parse CLI ===
+ args = parse_args()
-paths_img = sorted(glob.glob(config["paths_img"]))
-paths_img_nrg = sorted(glob.glob(config["paths_img_nrg"]))
-paths_mask = sorted(glob.glob(config["paths_mask"]))
-
-hue_min = float(config["hue_min"])
-hue_max = float(config["hue_max"])
-sat_thr = float(config["sat_thr"])
-val_thr = float(config["val_thr"])
-
-num_images = int(config["num_images"])
-num_compare = int(config["num_compare"])
-
-comparison_results = []
+ # === Load config from YAML ===
+ config = load_config(args.config)
-assert len(paths_img) > 0, "No RGB images found in config path"
-assert len(paths_img_nrg) > 0, "No NRG images found in config path"
-assert len(paths_mask) > 0, "No main masks found in config path"
+ # === CLI overrides config ===
-# Program Functions
+ # Output path
+ if args.output_dir:
+ config["output_dir"] = args.output_dir
-def generate_segment_mask_rgb(rgb_image, hue_min, hue_max, sat_thr, val_thr):
+ # Input paths
+ if args.path_rgb:
+ config["paths_img"] = args.path_rgb
+ if args.path_nrg:
+ config["paths_img_nrg"] = args.path_nrg
+ if args.path_mask:
+ config["paths_mask"] = args.path_mask
- hsv = color.rgb2hsv(rgb_image.astype(np.float32) / 255.0)
+ # Thresholds
+ if args.hue_min is not None:
+ config["hue_min"] = args.hue_min
+ if args.hue_max is not None:
+ config["hue_max"] = args.hue_max
+ if args.sat_thr is not None:
+ config["sat_thr"] = args.sat_thr
+ if args.val_thr is not None:
+ config["val_thr"] = args.val_thr
- H = hsv[:, :, 0]
- S = hsv[:, :, 1]
- V = hsv[:, :, 2]
-
- hue_mask = (H > hue_min) & (H < hue_max)
- saturation_mask = S > sat_thr
- value_mask = V > val_thr
-
- mask_rgb = hue_mask & saturation_mask & value_mask
- return mask_rgb.astype(np.uint8)
-
-def generate_segment_mask_nrg(nrg_image):
- nrg = nrg_image.astype(np.float32)
-
- nir = nrg[:,:,0]
- red = nrg[:,:,1]
- green = nrg[:,:,2]
-
- nir_norm = (nir - nir.min()) / (nir.max() - nir.min() + 1e-6)
- red_norm = (red - red.min()) / (red.max() - red.min() + 1e-6)
- green_norm = (green - green.min()) / (green.max() - green.min() + 1e-6)
-
- mask = (
- (nir_norm < 0.40) &
- (red_norm > 0.55) &
- (green_norm > 0.55)
+ # General
+ if args.num_images is not None:
+ config["num_images"] = args.num_images
+ if args.num_compare is not None:
+ config["num_compare"] = args.num_compare
+
+ # === Load image paths ===
+ rgb_paths = sorted(glob.glob(config["paths_img"]))
+ nrg_paths = sorted(glob.glob(config["paths_img_nrg"]))
+ mask_paths = sorted(glob.glob(config["paths_mask"]))
+
+ assert rgb_paths and nrg_paths and mask_paths, "No input files found"
+
+ if config["num_images"] is not None:
+ rgb_paths = rgb_paths[:config["num_images"]]
+ nrg_paths = nrg_paths[:config["num_images"]]
+ mask_paths = mask_paths[:config["num_images"]]
+
+ # === Number of images used for evaluation ===
+ n = min(
+ config["num_compare"],
+ len(rgb_paths),
+ len(nrg_paths),
+ len(mask_paths)
)
- return mask.astype(np.uint8)
-
-def connection_mask(rgb_image, nrg_image, hue_min, hue_max, sat_thr, val_thr):
-
- mask_rgb = generate_segment_mask_rgb(
- rgb_image, hue_min, hue_max, sat_thr, val_thr
- ).astype(np.uint8)
-
- mask_nrg = generate_segment_mask_nrg(nrg_image).astype(np.uint8)
-
- combined = np.zeros_like(mask_rgb, dtype=bool)
-
- combined[(mask_rgb == 1) & (mask_nrg == 1)] = 1
- combined[(mask_rgb == 1) & (mask_nrg == 0)] = 1
- combined[(mask_rgb == 0) & (mask_nrg == 1)] = 0
- combined[(mask_rgb == 0) & (mask_nrg == 0)] = 0
-
-
- combined = morphology.binary_opening(combined, morphology.disk(1))
- combined = morphology.binary_closing(combined, morphology.disk(2))
- combined = morphology.remove_small_objects(combined, min_size=50)
- combined = morphology.remove_small_holes(combined, area_threshold=100)
- combined = mask_rgb & morphology.binary_dilation(mask_nrg, morphology.disk(2))
- combined = morphology.binary_closing(combined, morphology.disk(4))
-
-
-
- return combined.astype(np.uint8)
-
-def combined_masks(img, nrg):
- mask_rgb = generate_segment_mask_rgb(
- img, hue_min, hue_max, sat_thr, val_thr
- ).astype(bool)
-
- mask_nrg = generate_segment_mask_nrg(nrg).astype(bool)
-
- combined = mask_rgb | mask_nrg
-
- return combined.astype(np.uint8)
-
-
-
-def iou(mask1, mask2):
- inter = np.logical_and(mask1, mask2)
- all = np.logical_or(mask1, mask2)
-
- inter_sum = np.sum(inter)
- all_sum = np.sum(all)
-
- iou_score = inter_sum / all_sum
- return iou_score
-
-def preview_loaded_images(num_images):
- if len(paths_img) == 0:
- print("No images found")
- return
- preview_paths = paths_img[:num_images]
- preview_paths_nrg = paths_img_nrg[:num_images]
- preview_paths_mask = paths_mask[:num_images]
-
- idx = 0
-
- fig, axes = plt.subplots(1, 3, figsize=(18, 6))
-
- def draw():
- img = io.imread(preview_paths[idx])
- nrg = io.imread(preview_paths_nrg[idx])
- mask = cv2.imread(preview_paths_mask[idx], cv2.IMREAD_GRAYSCALE)
-
- axes[0].imshow(img)
- axes[0].set_title("RGB Image")
-
- axes[1].imshow(nrg)
- axes[1].set_title("NRG Image")
-
- axes[2].imshow(mask, cmap="gray")
- axes[2].set_title("Ground Truth Mask")
-
- for ax in axes:
- ax.axis("off")
-
- fig.suptitle(
- f"{idx+1}/{len(preview_paths)} | {os.path.basename(preview_paths[idx])}",
- fontsize=14
- )
-
- fig.canvas.draw_idle()
-
- def on_key(event):
- nonlocal idx
- if event.key == "right":
- idx = (idx + 1) % len(preview_paths)
- draw()
- elif event.key == "left":
- idx = (idx - 1) % len(preview_paths)
- draw()
-
- fig.canvas.mpl_connect("key_press_event", on_key)
- draw()
- plt.show()
-
-
-def run_evaluation():
- comparison_results.clear()
-
- for i in range(min(num_compare, len(paths_img))):
- img_path = paths_img[i]
- nrg_path = paths_img_nrg[i]
- main_mask_path = paths_mask[i]
-
- img = io.imread(img_path)
- nrg = io.imread(nrg_path)
- generated_mask = connection_mask(img, nrg, hue_min, hue_max, sat_thr, val_thr);
- rgb_mask = generate_segment_mask_rgb(img, hue_min, hue_max, sat_thr, val_thr);
- nrg_mask = generate_segment_mask_nrg(nrg);
- combined_mask = combined_masks(img, nrg)
-
- main_mask = cv2.imread(main_mask_path, cv2.IMREAD_GRAYSCALE);
-
- iou_score = iou(generated_mask, main_mask);
-
- comparison_results.append({
- 'original_image': img,
- 'generated_mask': generated_mask,
- 'combined_mask': combined_mask,
- 'rgb_mask': rgb_mask,
- 'nrg_mask': nrg_mask,
- 'main_mask': main_mask,
- 'iou_score': iou_score,
- 'image_filename': os.path.basename(img_path)
- });
-
-
-def show_comparison_results():
- if len(comparison_results) == 0:
- print("No comparison results to show")
- return
-
- idx = 0
-
- fig, axes = plt.subplots(2, 3, figsize=(16, 10))
- axes = axes.flatten()
-
- def draw():
- result = comparison_results[idx]
-
- axes[0].imshow(result['original_image'])
- axes[0].set_title(result['image_filename'])
-
- axes[1].imshow(result['generated_mask'], cmap='gray')
- axes[1].set_title("Generated Mask")
-
- axes[2].imshow(result['main_mask'], cmap='gray')
- axes[2].set_title("Main Mask")
-
- axes[3].imshow(result['rgb_mask'], cmap='gray')
- axes[3].set_title("RGB Mask")
-
- axes[4].imshow(result['nrg_mask'], cmap='gray')
- axes[4].set_title("NRG Mask")
-
- axes[5].imshow(result['combined_mask'], cmap='gray')
- axes[5].set_title("Combined RGB + NRG")
-
- for ax in axes:
- ax.axis("off")
-
- fig.suptitle(
- f"{idx+1}/{len(comparison_results)} | IoU: {result['iou_score']*100:.1f}%",
- fontsize=14
+ results = []
+ # === Main processing loop ===
+ for i in range(n):
+ rgb = io.imread(rgb_paths[i])
+ nrg = io.imread(nrg_paths[i])
+ gt = cv2.imread(mask_paths[i], cv2.IMREAD_GRAYSCALE)
+
+ rgb_mask = generate_segment_mask_rgb(
+ rgb,
+ config["hue_min"],
+ config["hue_max"],
+ config["sat_thr"],
+ config["val_thr"]
)
- fig.canvas.draw_idle()
-
- def on_key(event):
- nonlocal idx
- if event.key == "right":
- idx = (idx + 1) % len(comparison_results)
- draw()
- elif event.key == "left":
- idx = (idx - 1) % len(comparison_results)
- draw()
-
- fig.canvas.mpl_connect("key_press_event", on_key)
- draw()
- plt.show()
-
-def result_all_stats():
- iou_scores = [result['iou_score'] for result in comparison_results]
- image_filenames = [os.path.basename(result['image_filename']) for result in comparison_results]
-
- plt.figure(figsize=(15, 7))
- plt.bar(image_filenames, iou_scores, color='skyblue')
- plt.xlabel('Nazwa pliku obrazu')
- plt.ylabel('Wynik IoU')
- plt.title('Comparison results IoUs for all sets')
- plt.xticks(rotation=45, ha='right')
- plt.ylim(0, 1)
- plt.tight_layout()
- plt.show()
-
-def confusion_matrix_all():
- y_main_all = []
- y_pred_all = []
-
- for result in comparison_results:
- main_mask = result['main_mask'].astype(bool).flatten()
- pred_mask = result['generated_mask'].astype(bool).flatten()
-
- y_main_all.extend(main_mask)
- y_pred_all.extend(pred_mask)
-
- cm = confusion_matrix(y_main_all, y_pred_all)
- print("CONFUSION MATRIX (for all images)")
- print("TN FP")
- print("FN TP")
- print(cm)
- return confusion_matrix(y_main_all, y_pred_all)
-
-def plot_confusion_matrix(cm, title="Confusion Matrix – normalized (%)"):
-
- labels = ["Background", "Dead Tree"]
-
- cm = cm.astype(np.float32)
- cm_percent = cm / cm.sum(axis=1, keepdims=True) * 100
-
- fig, ax = plt.subplots(figsize=(6, 5))
- im = ax.imshow(cm_percent, vmin=0, vmax=100)
-
- ax.set_title(title, fontsize=14, pad=10)
- ax.set_xlabel("Prediction", fontsize=12)
- ax.set_ylabel("Main Truth", fontsize=12)
-
- ax.set_xticks([0, 1])
- ax.set_yticks([0, 1])
- ax.set_xticklabels(labels)
- ax.set_yticklabels(labels)
-
- # Explicit cell labels
- cell_names = [
- ["TN\n(Background → Background)", "FP\n(Background → Dead Tree)"],
- ["FN\n(Dead Tree → Background)", "TP\n(Dead Tree → Dead Tree)"]
- ]
-
- for i in range(2):
- for j in range(2):
- value = cm_percent[i, j]
- text_color = "white" if value > 50 else "black"
-
- ax.text(
- j, i,
- f"{cell_names[i][j]}\n{value:.1f}%",
- ha="center",
- va="center",
- color=text_color,
- fontsize=11,
- fontweight="bold"
- )
-
- cbar = fig.colorbar(im, ax=ax)
- cbar.set_label("Percentage of pixels (%)")
-
- plt.tight_layout()
- plt.show()
-
+ nrg_mask = generate_segment_mask_nrg(nrg)
+ fused = fuse_masks(rgb_mask, nrg_mask)
+ score = iou(fused, gt)
-def search_best_thresholds(hue_ranges, sats, vals):
- rgb_images = [io.imread(p) for p in paths_img[:num_compare]]
- main_masks = [cv2.imread(p, cv2.IMREAD_GRAYSCALE).astype(bool)
- for p in paths_mask[:num_compare]]
+ results.append({
+ "image_filename": os.path.basename(rgb_paths[i]),
+ "generated_mask": fused,
+ "main_mask": gt,
+ "iou_score": score
+ })
- hsv_images = [color.rgb2hsv(img.astype(np.float32) / 255.0) for img in rgb_images]
+ # === Evaluation ===
+ logging.info(f"{n} images have been loaded to processing\n and {len(results)} images have been processed correctly")
+ cm = compute_confusion_matrix(results)
+ logging.info(f"Confusion Matrix:\n{cm}")
- best_iou = -1
- best_params = (None, None, None, None)
+ # === Save results ===
+ save_iou_results(results)
+ save_masks(results, config["output_dir"])
- hue_iou = []
- sat_iou = []
- val_iou = []
-
- for (hmin, hmax) in hue_ranges:
- ious = []
- for s in sats:
- for v in vals:
- for hsv, main in zip(hsv_images, main_masks):
- H, S, V = hsv[:,:,0], hsv[:,:,1], hsv[:,:,2]
- mask = (H > hmin) & (H < hmax) & (S > s) & (V > v)
-
- inter = np.logical_and(mask, main).sum()
- union = np.logical_or(mask, main).sum()
- ious.append(inter / union if union != 0 else 0)
-
- hue_iou.append(np.mean(ious))
-
- for s in sats:
- ious = []
- for (hmin, hmax) in hue_ranges:
- for v in vals:
- for hsv, main in zip(hsv_images, main_masks):
- H, S, V = hsv[:,:,0], hsv[:,:,1], hsv[:,:,2]
- mask = (H > hmin) & (H < hmax) & (S > s) & (V > v)
-
- inter = np.logical_and(mask, main).sum()
- union = np.logical_or(mask, main).sum()
- ious.append(inter / union if union != 0 else 0)
-
- sat_iou.append(np.mean(ious))
-
- for v in vals:
- ious = []
- for (hmin, hmax) in hue_ranges:
- for s in sats:
- for hsv, main in zip(hsv_images, main_masks):
- H, S, V = hsv[:,:,0], hsv[:,:,1], hsv[:,:,2]
- mask = (H > hmin) & (H < hmax) & (S > s) & (V > v)
-
- inter = np.logical_and(mask, main).sum()
- union = np.logical_or(mask, main).sum()
- ious.append(inter / union if union != 0 else 0)
-
- val_iou.append(np.mean(ious))
-
- for (hmin, hmax) in hue_ranges:
- for s in sats:
- for v in vals:
- ious = []
- for hsv, main in zip(hsv_images, main_masks):
- H, S, V = hsv[:,:,0], hsv[:,:,1], hsv[:,:,2]
- mask = (H > hmin) & (H < hmax) & (S > s) & (V > v)
-
- inter = np.logical_and(mask, main).sum()
- union = np.logical_or(mask, main).sum()
- ious.append(inter / union if union != 0 else 0)
-
- mean_iou = np.mean(ious)
- if mean_iou > best_iou:
- best_iou = mean_iou
- best_params = (hmin, hmax, s, v)
-
-
- plt.figure(figsize=(18,5))
-
- plt.subplot(1,3,1)
- plt.bar(range(len(hue_ranges)), hue_iou)
- plt.xticks(range(len(hue_ranges)),
- [f"{h[0]:.2f}-{h[1]:.2f}" for h in hue_ranges],
- rotation=45)
- plt.title("IoU vs Hue range")
- plt.ylabel("Mean IoU")
-
- plt.subplot(1,3,2)
- plt.bar(range(len(sats)), sat_iou)
- plt.xticks(range(len(sats)), [f"{s:.2f}" for s in sats])
- plt.title("IoU vs Saturation")
-
- plt.subplot(1,3,3)
- plt.bar(range(len(vals)), val_iou)
- plt.xticks(range(len(vals)), [f"{v:.2f}" for v in vals])
- plt.title("IoU vs Value")
-
- plt.tight_layout()
- plt.show()
-
- print("Best IoU:", best_iou)
- print("Best params:")
- print("Hue:", best_params[0], "-", best_params[1])
- print("S :", best_params[2])
- print("V :", best_params[3])
-
- return best_params, best_iou
-
-def save_iou_results(filename="iou_results.csv"):
- if len(comparison_results) == 0:
- print("No IoU results to save")
- return
-
- script_dir = os.path.dirname(os.path.abspath(__file__))
- save_path = os.path.join(script_dir, filename)
-
- with open(save_path, "w", encoding="utf-8") as f:
- f.write("image_name - iou\n")
- for result in comparison_results:
- f.write(f"{result['image_filename']} - {result['iou_score']*100:.2f}%\n")
-
- print(f"IoU results saved to: {save_path}")
-
-# Saves generated segmentation masks to a specified output directory
-def save_generated_masks(output_dir):
- if len(comparison_results) == 0:
- print("No generated masks to save")
- return
-
- os.makedirs(output_dir, exist_ok=True)
-
- for result in comparison_results:
- mask = (result["generated_mask"] * 255).astype(np.uint8)
- filename = os.path.splitext(result["image_filename"])[0] + "_mask.png"
- save_path = os.path.join(output_dir, filename)
-
- cv2.imwrite(save_path, mask)
-
- print(f"Generated masks saved to: {output_dir}")
-
-
-def main():
- output_dir = config["output_dir"]
-
- preview_loaded_images(num_images)
- run_evaluation()
- show_comparison_results()
- result_all_stats()
- cm = confusion_matrix_all()
- plot_confusion_matrix(cm)
- save_iou_results()
- save_generated_masks(output_dir)
+ logging.info("Pipeline finished successfully")
if __name__ == "__main__":
- main()
-
-search_best_thresholds(
- hue_ranges=[(0.7, 0.8), (0.75, 0.85), (0.8, 0.9)],
- sats=[0.15, 0.2],
- vals=[0.5, 0.55]
-)
+ main()
diff --git a/src/__pycache__/cli.cpython-313.pyc b/src/__pycache__/cli.cpython-313.pyc
new file mode 100644
index 0000000..7ed3df3
Binary files /dev/null and b/src/__pycache__/cli.cpython-313.pyc differ
diff --git a/src/__pycache__/config_loader.cpython-313.pyc b/src/__pycache__/config_loader.cpython-313.pyc
new file mode 100644
index 0000000..9cb2031
Binary files /dev/null and b/src/__pycache__/config_loader.cpython-313.pyc differ
diff --git a/src/__pycache__/evaluation.cpython-313.pyc b/src/__pycache__/evaluation.cpython-313.pyc
new file mode 100644
index 0000000..332eec0
Binary files /dev/null and b/src/__pycache__/evaluation.cpython-313.pyc differ
diff --git a/src/__pycache__/io_utils.cpython-313.pyc b/src/__pycache__/io_utils.cpython-313.pyc
new file mode 100644
index 0000000..177b953
Binary files /dev/null and b/src/__pycache__/io_utils.cpython-313.pyc differ
diff --git a/src/__pycache__/segmentation.cpython-313.pyc b/src/__pycache__/segmentation.cpython-313.pyc
new file mode 100644
index 0000000..ab6dee4
Binary files /dev/null and b/src/__pycache__/segmentation.cpython-313.pyc differ
diff --git a/src/cli.py b/src/cli.py
new file mode 100644
index 0000000..a6f7f19
--- /dev/null
+++ b/src/cli.py
@@ -0,0 +1,85 @@
+import argparse
+
+def parse_args():
+ parser = argparse.ArgumentParser(
+ description=(
+ "Dead Tree Segmentation Pipeline (RGB + NRG)\n\n"
+ "The pipeline performs classical image segmentation using RGB and NRG data,\n"
+ "computes evaluation metrics (IoU, confusion matrix) and saves results to disk."
+ ),
+ formatter_class=argparse.RawTextHelpFormatter
+ )
+
+ # === Configuration ===
+ parser.add_argument(
+ "-c", "--config",
+ type=str,
+ default="config.yaml",
+ help="Path to YAML configuration file (default: config.yaml)"
+ )
+
+ # === Segmentation thresholds ===
+ parser.add_argument(
+ "-h_min", "--hue_min",
+ type=float,
+ help="Minimum HSV hue threshold for RGB segmentation"
+ )
+ parser.add_argument(
+ "-h_max", "--hue_max",
+ type=float,
+ help="Maximum HSV hue threshold for RGB segmentation"
+ )
+ parser.add_argument(
+ "-s", "--sat_thr",
+ type=float,
+ help="HSV saturation threshold for RGB segmentation"
+ )
+ parser.add_argument(
+ "-v", "--val_thr",
+ type=float,
+ help="HSV value (brightness) threshold for RGB segmentation"
+ )
+
+ # === General pipeline control ===
+ parser.add_argument(
+ "-n_img", "--num_images",
+ type=int,
+ help="Number of images to load or preview from the dataset"
+ )
+ parser.add_argument(
+ "-n_cmp", "--num_compare",
+ type=int,
+ help="Number of images used for quantitative evaluation"
+ )
+
+ # === Input / Output paths ===
+ parser.add_argument(
+ "-o", "--output_dir",
+ type=str,
+ help="Output directory for generated masks, metrics and logs"
+ )
+ parser.add_argument(
+ "-p_rgb", "--path_rgb",
+ type=str,
+ help="Path or glob pattern to RGB images (e.g. data/RGB_images/*.png)"
+ )
+ parser.add_argument(
+ "-p_nrg", "--path_nrg",
+ type=str,
+ help="Path or glob pattern to NRG images (e.g. data/NRG_images/*.png)"
+ )
+ parser.add_argument(
+ "-p_mask", "--path_mask",
+ type=str,
+ help="Path or glob pattern to ground truth masks (e.g. data/masks/*.png)"
+ )
+
+ # === Logging ===
+ parser.add_argument(
+ "--log-level",
+ default="INFO",
+ choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+ help="Logging verbosity level (default: INFO)"
+ )
+
+ return parser.parse_args()
diff --git a/src/config_loader.py b/src/config_loader.py
new file mode 100644
index 0000000..0a74132
--- /dev/null
+++ b/src/config_loader.py
@@ -0,0 +1,23 @@
+import yaml
+
+def load_config(path="config.yaml"):
+ with open(path, "r", encoding="utf-8") as f:
+ cfg = yaml.safe_load(f)
+
+ config = {
+ "paths_img": cfg["paths"]["rgb"],
+ "paths_img_nrg": cfg["paths"]["nrg"],
+ "paths_mask": cfg["paths"]["masks"],
+ "output_dir": cfg["paths"]["output"],
+ "hue_min": float(cfg["thresholds"]["hue_min"]),
+ "hue_max": float(cfg["thresholds"]["hue_max"]),
+ "sat_thr": float(cfg["thresholds"]["sat_thr"]),
+ "val_thr": float(cfg["thresholds"]["val_thr"]),
+ "num_images": int(cfg["general"]["num_images"]),
+ "num_compare": int(cfg["general"]["num_compare"]),
+ }
+
+ if config["num_compare"] > config["num_images"]:
+ raise ValueError("num_compare cannot exceed num_images")
+
+ return config
diff --git a/src/evaluation.py b/src/evaluation.py
new file mode 100644
index 0000000..e99f6f8
--- /dev/null
+++ b/src/evaluation.py
@@ -0,0 +1,20 @@
+import numpy as np
+import logging
+from sklearn.metrics import confusion_matrix
+
+logger = logging.getLogger("dead_tree")
+
+def iou(mask1, mask2):
+ inter = np.logical_and(mask1, mask2).sum()
+ union = np.logical_or(mask1, mask2).sum()
+ return inter / union if union > 0 else 0
+
+def compute_confusion_matrix(results):
+ y_true, y_pred = [], []
+
+ for r in results:
+ y_true.extend(r["main_mask"].astype(bool).flatten())
+ y_pred.extend(r["generated_mask"].astype(bool).flatten())
+
+ logger.info("Confusion matrix computed")
+ return confusion_matrix(y_true, y_pred)
diff --git a/src/io_utils.py b/src/io_utils.py
new file mode 100644
index 0000000..c46f1a9
--- /dev/null
+++ b/src/io_utils.py
@@ -0,0 +1,23 @@
+import os, cv2, logging, csv
+
+logger = logging.getLogger("dead_tree")
+
+def save_iou_results(results, filename="iou_results.csv"):
+ with open(filename, "w", encoding="utf-8") as f:
+ f.write("image_name,iou\n")
+ for r in results:
+ f.write(f"{r['image_filename']},{r['iou_score']*100:.2f}%\n")
+
+ logger.info(f"IoU results saved to {filename}")
+
+def save_masks(results, output_dir):
+ os.makedirs(output_dir, exist_ok=True)
+
+ for r in results:
+ path = os.path.join(
+ output_dir,
+ os.path.splitext(r["image_filename"])[0] + "_mask.png"
+ )
+ cv2.imwrite(path, r["generated_mask"] * 255)
+
+ logger.info(f"Masks saved to {output_dir}")
diff --git a/src/logger.py b/src/logger.py
new file mode 100644
index 0000000..0b0c771
--- /dev/null
+++ b/src/logger.py
@@ -0,0 +1,29 @@
+import logging
+import sys
+
+def setup_logger(
+ name: str = "dead_tree",
+ level: int = logging.INFO,
+ log_file: str | None = None
+) -> logging.Logger:
+ logger = logging.getLogger(name)
+ logger.setLevel(level)
+
+ if logger.handlers:
+ return logger
+
+ formatter = logging.Formatter(
+ "[%(asctime)s] [%(levelname)s] %(name)s: %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S"
+ )
+
+ console_handler = logging.StreamHandler(sys.stdout)
+ console_handler.setFormatter(formatter)
+ logger.addHandler(console_handler)
+
+ if log_file:
+ file_handler = logging.FileHandler(log_file, encoding="utf-8")
+ file_handler.setFormatter(formatter)
+ logger.addHandler(file_handler)
+
+ return logger
diff --git a/src/optimization.py b/src/optimization.py
new file mode 100644
index 0000000..7bd82b7
--- /dev/null
+++ b/src/optimization.py
@@ -0,0 +1,21 @@
+import numpy as np
+from skimage import color
+
+def search_best_thresholds(hsv_images, masks, hue_ranges, sats, vals):
+ best_iou, best_params = -1, None
+
+ for hmin, hmax in hue_ranges:
+ for s in sats:
+ for v in vals:
+ scores = []
+ for hsv, gt in zip(hsv_images, masks):
+ H,S,V = hsv[:,:,0], hsv[:,:,1], hsv[:,:,2]
+ mask = (H>hmin)&(Hs)&(V>v)
+ inter = (mask & gt).sum()
+ union = (mask | gt).sum()
+ scores.append(inter/union if union else 0)
+ mean_iou = np.mean(scores)
+ if mean_iou > best_iou:
+ best_iou, best_params = mean_iou, (hmin,hmax,s,v)
+
+ return best_params, best_iou
diff --git a/src/segmentation.py b/src/segmentation.py
new file mode 100644
index 0000000..aa26f8f
--- /dev/null
+++ b/src/segmentation.py
@@ -0,0 +1,22 @@
+import numpy as np
+from skimage import color, morphology
+
+def generate_segment_mask_rgb(rgb_image, hue_min, hue_max, sat_thr, val_thr):
+ hsv = color.rgb2hsv(rgb_image.astype(np.float32) / 255.0)
+ H, S, V = hsv[:,:,0], hsv[:,:,1], hsv[:,:,2]
+ return ((H > hue_min) & (H < hue_max) & (S > sat_thr) & (V > val_thr)).astype(np.uint8)
+
+def generate_segment_mask_nrg(nrg_image):
+ nrg = nrg_image.astype(np.float32)
+ nir, red, green = nrg[:,:,0], nrg[:,:,1], nrg[:,:,2]
+
+ nir = (nir - nir.min()) / (nir.max() - nir.min() + 1e-6)
+ red = (red - red.min()) / (red.max() - red.min() + 1e-6)
+ green = (green - green.min()) / (green.max() - green.min() + 1e-6)
+
+ return ((nir < 0.4) & (red > 0.55) & (green > 0.55)).astype(np.uint8)
+
+def fuse_masks(mask_rgb, mask_nrg):
+ combined = mask_rgb & morphology.binary_dilation(mask_nrg, morphology.disk(2))
+ combined = morphology.binary_closing(combined, morphology.disk(4))
+ return combined.astype(np.uint8)
diff --git a/temp_config.yaml b/temp_config.yaml
index cd17f5a..332239d 100644
--- a/temp_config.yaml
+++ b/temp_config.yaml
@@ -2,7 +2,7 @@ paths:
rgb: ./data/RGB_images/*.png #
nrg: ./data/NRG_images/*.png #
masks: ./data/masks/*.png #
- output: ./data/generated_masks/ #
+ output: ./data_examples/example_generated_mask/ #
thresholds:
hue_min: 0.75