wound_segmentation.main module

Main entry point for performing wound segmentation using a trained deep learning model.

This script provides a command-line interface (CLI) for running segmentation predictions on either a single image or a directory of images. It performs preprocessing, inference, and post-processing (mask overlay and saving results).

The script automatically downloads pretrained model weights from Google Drive if they are not found locally.

Functions

run_single_image(model, image_path, output_dir, threshold)

Runs segmentation on a single image and saves the result.

run_batch(model, input_dir, output_dir, threshold)

Runs segmentation on all valid images in a folder and saves the results.

main()

Parses command-line arguments, loads the model, and executes the appropriate pipeline.

Examples

Run inference on a single image: >>> python main.py --image sample.jpg

Run inference on a folder of images: >>> python main.py --input_dir data/images/

Notes

  • The model architecture is defined in 'model.py'.

  • Preprocessing and mask saving utilities are used from 'preprocessing.py' and 'utils.py'.

  • Weights are auto-downloaded from a public Google Drive link configured in 'constants.py'.

wound_segmentation.main.run_single_image(model: tensorflow.keras.Model, image_path: str, output_dir: str, threshold: float) → None

Run segmentation on a single input image.

This function runs the segmentation pipeline on a single image given as an input by the user. It calls the predict_mask function to predict the binary mask and saves the result to the given (or default) output directory.

Parameters:

  • model (tf.keras.Model) -- The loaded segmentation model.

  • image_path (str) -- Path to the input RGB image file.

  • output_dir (str) -- Directory where the predicted result will be saved.

  • threshold (float) -- Threshold for binarizing the model's predicted probability map.

Raises:

FileNotFoundError -- If the input image file does not exist.

Examples

>>> model = load_segmentation_model("model_weights/seg_model.weights.h5")
>>> run_single_image(model, "images/sample.jpg", "outputs/", threshold=0.5)
wound_segmentation.main.run_batch(model: tensorflow.keras.Model, input_dir: str, output_dir: str, threshold: float) → None

Run segmentation inference on all images in a given folder.

This function runs the segmentation pipeline on a batch of images given as an input directory by the user. It calls the predict_mask function to predict the binary mask and saves the result to the given (or default) output directory.

Parameters:

  • model (tf.keras.Model) -- The loaded segmentation model.

  • input_dir (str) -- Folder containing images to segment.

  • output_dir (str) -- Directory where the predicted results will be saved.

  • threshold (float) -- Threshold for binarizing the model's predicted probability map.

Raises:

FileNotFoundError -- If the input directory is missing or contains no valid images.

Examples

>>> run_batch(model, "dataset/images", "outputs", threshold=0.6)
wound_segmentation.main.main() → None

Parse command-line arguments and execute segmentation prediction.

This function Supports single image mode (--image) and batch mode (--input_dir). It also automatically downloads model weights if not available locally.

Raises:

SystemExit -- If CLI arguments are invalid or missing.

Examples

>>> python main.py --image image.jpg
>>> python main.py --input_dir images/ --threshold 0.6