Merge pull request #9 from iteal/develop

Develop

Author: iteal

README.md

@@ -2,7 +2,7 @@
 
 [![PyPI status](https://img.shields.io/pypi/status/wormpose.svg)](https://pypi.python.org/pypi/wormpose/)
 [![PyPI version](https://badge.fury.io/py/wormpose.svg)](https://pypi.python.org/pypi/wormpose/)
-[![Build Status](https://travis-ci.org/iteal/wormpose.svg?branch=main)](https://travis-ci.org/iteal/wormpose)
+[![Build Status](https://travis-ci.com/iteal/wormpose.svg?branch=main)](https://travis-ci.com/iteal/wormpose)
 [![PyPI pyversions](https://img.shields.io/pypi/pyversions/wormpose.svg)](https://pypi.python.org/pypi/wormpose/)
 [![DOI](https://zenodo.org/badge/246706178.svg)](https://zenodo.org/badge/latestdoi/246706178)
 [![Maintainability](https://img.shields.io/codeclimate/maintainability/iteal/wormpose)](https://codeclimate.com/github/iteal/wormpose/maintainability)

docs/source/install.rst

@@ -37,9 +37,15 @@ or directly call the Python functions.
    wormpose COMMAND -h
 
 
-
 **Troubleshooting**
 
 If the `wormpose` command is not found, maybe you installed WormPose with the pip option ``--user``. Please check
 that the local user folder where pip installs the script is in your PATH
-variable.
+variable.
+
+
+**Limitations**
+
+WormPose doesn't support at the moment multiple worms in one image.
+Only grayscale images (1 channel) are supported.
+The resolution of the worm should not be too low, we have tested it successfully with images where the worm length was 90 pixels, but having a smaller resolution will degrade the accuracy rapidly.

docs/source/tierpsy.rst

@@ -27,7 +27,8 @@ For using `Open Worm Movement Database <http://movement.openworm.org/>`__ videos
 
 **Limitations**
 
-For Tierpsy files with several worms per file, WormPose will only load one worm, the one with the smallest index.
+Please use the option "extract timestamp" in Tierpsy so that the field "timestamp/raw" is set in the hdf5 file. The configuration without the timestamp is not currently supported in WormPose.
+For Tierpsy files with several worm indexes per file, WormPose will only load one worm, the one with the smallest index.
 
 **Troubleshooting**
 

wormpose/commands/calibrate_dataset.py

@@ -16,6 +16,7 @@
 from wormpose.commands import _log_parameters
 from wormpose.config import default_paths
 from wormpose.dataset import Dataset
+from wormpose.dataset.image_processing.options import add_image_processing_arguments
 from wormpose.dataset.loader import get_dataset_name
 from wormpose.dataset.loader import load_dataset
 from wormpose.images.scoring.centerline_accuracy_check import CenterlineAccuracyCheck
@@ -155,7 +156,10 @@ def calibrate(dataset_loader: str, dataset_path: str, **kwargs):
     os.makedirs(calibration_results_dir, exist_ok=True)
 
     dataset = load_dataset(
-        dataset_loader, dataset_path, resize_options=args.resize_options, selected_video_names=args.video_names
+        dataset_loader,
+        dataset_path,
+        selected_video_names=args.video_names,
+        **vars(args),
     )
 
     calibrator = _Calibrator(
@@ -178,10 +182,8 @@ def main():
 
     parser = argparse.ArgumentParser()
 
-    # inputs
     parser.add_argument("dataset_loader", type=str)
     parser.add_argument("dataset_path", type=str)
-    add_resizing_arguments(parser)
     parser.add_argument(
         "--video_names",
         type=str,
@@ -199,6 +201,8 @@ def main():
         "--save_images", default=False, action="store_true", help="Also save the images used for the calibration"
     )
     parser.add_argument("--random_seed", type=int, help="Optional random seed for deterministic results")
+    add_resizing_arguments(parser)
+    add_image_processing_arguments(parser)
 
     args = parser.parse_args()
 

wormpose/commands/evaluate_model.py

@@ -23,6 +23,7 @@
 from wormpose.config.default_paths import CONFIG_FILENAME
 from wormpose.config.experiment_config import load_config, add_config_argument
 from wormpose.dataset import Dataset
+from wormpose.dataset.image_processing.options import WORM_IS_LIGHTER
 from wormpose.dataset.loader import load_dataset, get_dataset_name
 from wormpose.dataset.loaders.resizer import ResizeOptions
 from wormpose.images.scoring import ResultsScoring, BaseScoringDataManager
@@ -172,6 +173,7 @@ def evaluate(dataset_path: str, **kwargs):
         dataset_path=dataset_path,
         selected_video_names=args.video_names,
         resize_options=ResizeOptions(resize_factor=config.resize_factor),
+        **{WORM_IS_LIGHTER: config.worm_is_lighter},
     )
 
     pkl_filenames = _generate_synthetic_data(

wormpose/commands/generate_training_data.py

@@ -23,6 +23,10 @@
     CONFIG_FILENAME,
 )
 from wormpose.config.experiment_config import save_config, ExperimentConfig
+from wormpose.dataset.image_processing.options import (
+    add_image_processing_arguments,
+    WORM_IS_LIGHTER,
+)
 from wormpose.dataset.loader import get_dataset_name
 from wormpose.dataset.loader import load_dataset
 from wormpose.dataset.loaders.resizer import add_resizing_arguments, ResizeOptions
@@ -53,6 +57,8 @@ def _parse_arguments(kwargs: dict):
         kwargs["video_names"] = None
     if kwargs.get("random_seed") is None:
         kwargs["random_seed"] = None
+    if kwargs.get(WORM_IS_LIGHTER) is None:
+        kwargs[WORM_IS_LIGHTER] = False
     kwargs["temp_dir"] = tempfile.mkdtemp(dir=kwargs["temp_dir"])
     kwargs["resize_options"] = ResizeOptions(**kwargs)
 
@@ -89,8 +95,8 @@ def generate(dataset_loader: str, dataset_path: str, **kwargs):
     dataset = load_dataset(
         dataset_loader=dataset_loader,
         dataset_path=dataset_path,
-        resize_options=args.resize_options,
         selected_video_names=args.video_names,
+        **vars(args),
     )
 
     start = time.time()
@@ -129,6 +135,7 @@ def generate(dataset_loader: str, dataset_path: str, **kwargs):
             num_eval_samples=num_eval_samples,
             resize_factor=args.resize_options.resize_factor,
             video_names=dataset.video_names,
+            worm_is_lighter=getattr(args, WORM_IS_LIGHTER),
         ),
         os.path.join(experiment_dir, CONFIG_FILENAME),
     )
@@ -157,6 +164,7 @@ def main():
     parser.add_argument("--num_process", type=int, help="How many worker processes")
     parser.add_argument("--random_seed", type=int, help="Optional random seed for deterministic results")
     add_resizing_arguments(parser)
+    add_image_processing_arguments(parser)
     args = parser.parse_args()
 
     last_progress = None

wormpose/commands/predict_dataset.py

@@ -24,6 +24,7 @@
 from wormpose.config.default_paths import RESULTS_FILENAME, CONFIG_FILENAME
 from wormpose.config.experiment_config import load_config, add_config_argument
 from wormpose.dataset.features import Features
+from wormpose.dataset.image_processing.options import WORM_IS_LIGHTER
 from wormpose.dataset.loader import get_dataset_name
 from wormpose.dataset.loader import load_dataset
 from wormpose.dataset.loaders.resizer import ResizeOptions
@@ -171,6 +172,7 @@ def predict(dataset_path: str, **kwargs):
         dataset_path=dataset_path,
         selected_video_names=args.video_names,
         resize_options=ResizeOptions(resize_factor=config.resize_factor),
+        **{WORM_IS_LIGHTER: config.worm_is_lighter},
     )
 
     keras_model = tf.keras.models.load_model(args.model_path, compile=False)

wormpose/config/experiment_config.py

@@ -34,6 +34,7 @@ def __init__(
         theta_dimensions: int = None,
         resize_factor: float = None,
         video_names: List[str] = None,
+        worm_is_lighter: bool = None,
     ):
         self.num_train_samples = num_train_samples
         self.num_eval_samples = num_eval_samples
@@ -42,6 +43,7 @@ def __init__(
         self.theta_dimensions = theta_dimensions
         self.resize_factor = resize_factor
         self.video_names = video_names
+        self.worm_is_lighter = worm_is_lighter
 
 
 def save_config(experiment_config: ExperimentConfig, config_filepath: str):

wormpose/dataset/base_dataset.py

@@ -13,6 +13,9 @@ class BaseFramePreprocessing(ABC):
     This object must be pickable (no inner functions for example)
     """
 
+    def __init__(self, *args, **kwargs):
+        pass
+
     @abstractmethod
     def process(self, frame: np.ndarray) -> Tuple[np.ndarray, int]:
         """

wormpose/dataset/image_processing/image_utils.py

@@ -55,14 +55,17 @@ def segment_foreground(
     foreground_close_struct_element,
     foreground_dilate_struct_element,
     threshold_fn: Callable[[np.ndarray], int],
+    is_foreground_lighter_than_background: bool,
 ):
     """
     Processes a frame to isolate the object of interest (worm) from the background
 
     :param frame: image to process
     :param foreground_close_struct_element: morphological element to close holes in the foreground mask
     :param foreground_dilate_struct_element: morphological element to expand the foreground mask
-    :param threshold_fn: function that will return the threshold to separate forefround from background in a frame
+    :param threshold_fn: function that will return the threshold to separate foreground from background in a frame
+    :param is_foreground_lighter_than_background: set to True if the foreground object of interest is lighter
+     (pixel values are on average higher) than the background
     :return: segmentation mask with values of 1 for the worm object and 0 for the background,
         and average value of the background pixels
     """
@@ -73,10 +76,14 @@ def segment_foreground(
     # use the threshold to deduce background and foreground masks, fill in holes
     foreground_mask = (frame > 0).astype(np.uint8) * (frame < background_threshold).astype(np.uint8)
     foreground_mask = cv2.morphologyEx(foreground_mask, cv2.MORPH_CLOSE, foreground_close_struct_element)
-    background_mask = ((frame > 0).astype(np.uint8) - foreground_mask) > 0
+    background_mask = (((frame > 0).astype(np.uint8) - foreground_mask) > 0).astype(np.uint8)
+
+    # invert foreground and background masks if the foreground is lighter than the background
+    if is_foreground_lighter_than_background:
+        foreground_mask, background_mask = background_mask, foreground_mask
 
     # calculate the average background color
-    background_values = frame[background_mask]
+    background_values = frame[background_mask.astype(bool)]
     background_color = int(np.mean(background_values)) if len(background_values) > 0 else 0
     background_color = frame.dtype.type(background_color)