Skip to content

Commit

Permalink
Automatic Grayscale Encoding (#21)
Browse files Browse the repository at this point in the history
kozlov721 authored Jul 27, 2024
1 parent 895499f commit f872f94
Showing 1 changed file with 17 additions and 0 deletions.
17 changes: 17 additions & 0 deletions modelconverter/utils/config.py
Original file line number Diff line number Diff line change
@@ -169,6 +169,23 @@ def _validate_reverse_input_channels(self) -> Self:

return self

@model_validator(mode="after")
def _validate_grayscale_inputs(self) -> Self:
if self.layout is None:
return self

if "C" not in self.layout:
return self

assert self.shape is not None

channels = self.shape[self.layout.index("C")]
if channels == 1:
logger.info("Detected grayscale input. Setting encoding to GRAY.")
self.encoding.from_ = self.encoding.to = Encoding.GRAY

return self

@model_validator(mode="before")
@classmethod
def _random_calibration(cls, data: Dict[str, Any]) -> Dict[str, Any]:

0 comments on commit f872f94

Please sign in to comment.