Skip to content

Commit

Permalink
Added new fields to accept weights only
Browse files Browse the repository at this point in the history
  • Loading branch information
IgorTatarnikov committed Oct 7, 2024
1 parent 594b5c0 commit 843da73
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 0 deletions.
14 changes: 14 additions & 0 deletions cellfinder/napari/detect/detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,9 @@ def widget(
classification_options,
skip_classification: bool,
use_pre_trained_weights: bool,
weights_only: bool,
trained_model: Optional[Path],
model_weights: Optional[Path],
batch_size: int,
misc_options,
start_plane: int,
Expand Down Expand Up @@ -299,6 +301,8 @@ def widget(
should be attempted
use_pre_trained_weights : bool
Select to use pre-trained model weights
weights_only : bool
Select to only provide the model weights
batch_size : int
How many points to classify at one time
skip_classification : bool
Expand All @@ -307,6 +311,9 @@ def widget(
trained_model : Optional[Path]
Trained model file path (home directory (default) -> pretrained
weights)
model_weights : Optional[Path]
Model weights file path (home directory (default) -> pretrained
weights)
start_plane : int
First plane to process (to process a subset of the data)
end_plane : int
Expand Down Expand Up @@ -372,12 +379,19 @@ def widget(
max_cluster_size,
)

if weights_only:
trained_model = None

if use_pre_trained_weights:
trained_model = None
model_weights = None

classification_inputs = ClassificationInputs(
skip_classification,
use_pre_trained_weights,
weights_only,
trained_model,
model_weights,
batch_size,
)

Expand Down
5 changes: 5 additions & 0 deletions cellfinder/napari/detect/detect_containers.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,12 +113,15 @@ class ClassificationInputs(InputContainer):

skip_classification: bool = False
use_pre_trained_weights: bool = True
weights_only: bool = False
trained_model: Optional[Path] = Path.home()
model_weights: Optional[Path] = Path.home()
batch_size: int = 64

def as_core_arguments(self) -> dict:
args = super().as_core_arguments()
del args["use_pre_trained_weights"]
del args["weights_only"]
return args

@classmethod
Expand All @@ -128,7 +131,9 @@ def widget_representation(cls) -> dict:
use_pre_trained_weights=dict(
value=cls.defaults()["use_pre_trained_weights"]
),
weights_only=dict(value=cls.defaults()["weights_only"]),
trained_model=dict(value=cls.defaults()["trained_model"]),
model_weights=dict(value=cls.defaults()["model_weights"]),
skip_classification=dict(
value=cls.defaults()["skip_classification"]
),
Expand Down

0 comments on commit 843da73

Please sign in to comment.