Skip to content

Commit

Permalink
fix(nrtk_transforms): avoid runtime error from type hints
Browse files Browse the repository at this point in the history
  • Loading branch information
PaulHax committed Aug 21, 2024
1 parent bd67b41 commit 309b817
Showing 1 changed file with 11 additions and 6 deletions.
17 changes: 11 additions & 6 deletions src/nrtk_explorer/library/nrtk_transforms.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any, Dict, Union
from typing import Any, Dict, Optional, TYPE_CHECKING

import numpy as np
import logging
Expand All @@ -15,22 +15,27 @@
from pybsm.otf import darkCurrentFromDensity
from nrtk.impls.perturb_image.generic.cv2.blur import GaussianBlurPerturber
from nrtk.impls.perturb_image.pybsm.perturber import PybsmPerturber, PybsmSensor, PybsmScenario

GaussianBlurPerturberType = Union[GaussianBlurPerturber, None]
PybsmPerturberType = Union[PybsmPerturber, None]
except ImportError:
logger.info("Disabling NRTK transforms due to missing library/failing imports")
ENABLED_NRTK_TRANSFORMS = False

if TYPE_CHECKING:
GaussianBlurPerturberType = GaussianBlurPerturber
PybsmPerturberType = PybsmPerturber
else:
GaussianBlurPerturberType = None
PybsmPerturberType = None

GaussianBlurPerturberArg = Optional[GaussianBlurPerturberType]
PybsmPerturberArg = Optional[PybsmPerturberType]


def nrtk_transforms_available():
return ENABLED_NRTK_TRANSFORMS


class NrtkGaussianBlurTransform(ImageTransform):
def __init__(self, perturber: GaussianBlurPerturberType = None):
def __init__(self, perturber: GaussianBlurPerturberArg = None):
if perturber is None:
perturber = GaussianBlurPerturber()

Expand Down Expand Up @@ -160,7 +165,7 @@ def createSampleSensorAndScenario():


class NrtkPybsmTransform(ImageTransform):
def __init__(self, perturber: PybsmPerturberType = None):
def __init__(self, perturber: PybsmPerturberArg = None):
if perturber is None:
sensor, scenario = createSampleSensorAndScenario()
perturber = PybsmPerturber(sensor=sensor, scenario=scenario)
Expand Down

0 comments on commit 309b817

Please sign in to comment.