# Path: common/loader.py
# def load_defense_and_attack(defense_path, attack_name, torch_model):
#     """Load modules with defense and attack.
# 
#     Args:
#         defense_path: path to directory with defense,
#             defense code expected to be in model.py file in this directory.
#         attack_name: name of the attack file in the defense directory.
# 
#     Returns:
#         defense_model: instance of defense model.
#         attack_cls: python class with attack.
#         task_def: task definition.
#         dataset_name: name of the dataset to use.
#     """
#     # Load modules
#     task_def_module = load_module('defense.task_definition',
#                                   os.path.join(defense_path, 'task_definition.py'))
#     defense_model_module = load_module('defense.model',
#                                        os.path.join(defense_path, 'model.py'))
#     attack_module = load_module('defense.attack',
#                                 os.path.join(defense_path, attack_name))
# 
#     # Finding task definition
#     if attack_name not in task_def_module.TASKS:
#         print('Task definition not found for attack {0}'.format(attack_name))
#         exit(1)
#     task_def = task_def_module.TASKS[attack_name]
# 
#     # Loading defense model
#     defense_classes = find_subclasses_in_module(defense_model_module, DefenseModel)
#     if torch_model:
#         defense_classes = [x for x in defense_classes if 'Torch' in x.__name__]
#     else:
#         defense_classes = [x for x in defense_classes if 'Torch' not in x.__name__]
#     if not defense_classes or len(defense_classes) > 1:
#         print('Defense {0} must have exactly one class implementing DefenseModel'.format(defense_path))
#         print('Found: ', defense_classes)
#         exit(1)
#     defense_model = defense_classes[0]()
# 
#     # Loading attack
#     attack_classes = find_subclasses_in_module(attack_module, Attack)
#     if not attack_classes or len(attack_classes) > 1:
#         print('Attack {0} must have exactly one class implementing Attack'.format(
#             os.path.join(defense_path, attack_name)))
#         print('Found: ', attack_classes)
#         exit(1)
#     attack_cls = attack_classes[0]
# 
#     return defense_model, attack_cls, task_def, task_def_module.DATASET


from typing import Callable, Optional
from art.attacks.attack import EvasionAttack
from art.estimators.classification import TensorFlowV2Classifier
from common.loader import load_defense_and_attack
import inspect
import os
import numpy as np
import tensorflow as tf
import armory.data.datasets
        self._model = defense_model
        self._attack_cls = attack_cls
        self._task_def = task_def

    def generate(self, x: np.ndarray, y: Optional[np.ndarray] = None) -> np.ndarray:
        assert y is not None, 'Labels has to be provided.'
        return self._attack_cls(self._task_def).attack(self._model, x, y)


class DefenseWrapper(TensorFlowV2Classifier):
    """ART wrapper for defenses."""

    def __init__(self, model):
        super().__init__(
            lambda x, **kwargs: model.classify(x),
            nb_classes=NUM_CLASSES,
            input_shape=INPUT_SHAPE,
            clip_values=(0., 1.),
            loss_object=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))

    def set_learning_phase(self, train):
        # Without overloading this method, armory does not work
        pass


def get_art_model(model_kwargs, wrapper_kwargs, weights_file):
    del wrapper_kwargs
    del weights_file
    defense_name = model_kwargs['defense']
    defense_path = get_defense_path(defense_name)

# Path: common/loader.py
# def load_defense_and_attack(defense_path, attack_name, torch_model):
#     """Load modules with defense and attack.
# 
#     Args:
#         defense_path: path to directory with defense,
#             defense code expected to be in model.py file in this directory.
#         attack_name: name of the attack file in the defense directory.
# 
#     Returns:
#         defense_model: instance of defense model.
#         attack_cls: python class with attack.
#         task_def: task definition.
#         dataset_name: name of the dataset to use.
#     """
#     # Load modules
#     task_def_module = load_module('defense.task_definition',
#                                   os.path.join(defense_path, 'task_definition.py'))
#     defense_model_module = load_module('defense.model',
#                                        os.path.join(defense_path, 'model.py'))
#     attack_module = load_module('defense.attack',
#                                 os.path.join(defense_path, attack_name))
# 
#     # Finding task definition
#     if attack_name not in task_def_module.TASKS:
#         print('Task definition not found for attack {0}'.format(attack_name))
#         exit(1)
#     task_def = task_def_module.TASKS[attack_name]
# 
#     # Loading defense model
#     defense_classes = find_subclasses_in_module(defense_model_module, DefenseModel)
#     if torch_model:
#         defense_classes = [x for x in defense_classes if 'Torch' in x.__name__]
#     else:
#         defense_classes = [x for x in defense_classes if 'Torch' not in x.__name__]
#     if not defense_classes or len(defense_classes) > 1:
#         print('Defense {0} must have exactly one class implementing DefenseModel'.format(defense_path))
#         print('Found: ', defense_classes)
#         exit(1)
#     defense_model = defense_classes[0]()
# 
#     # Loading attack
#     attack_classes = find_subclasses_in_module(attack_module, Attack)
#     if not attack_classes or len(attack_classes) > 1:
#         print('Attack {0} must have exactly one class implementing Attack'.format(
#             os.path.join(defense_path, attack_name)))
#         print('Found: ', attack_classes)
#         exit(1)
#     attack_cls = attack_classes[0]
# 
#     return defense_model, attack_cls, task_def, task_def_module.DATASET


# Path: armory_compat/wrappers.py
from typing import Callable, Optional
from art.attacks.attack import EvasionAttack
from art.estimators.classification import TensorFlowV2Classifier
from common.loader import load_defense_and_attack
import inspect
import os
import numpy as np
import tensorflow as tf
import armory.data.datasets

        self._model = defense_model
        self._attack_cls = attack_cls
        self._task_def = task_def

    def generate(self, x: np.ndarray, y: Optional[np.ndarray] = None) -> np.ndarray:
        assert y is not None, 'Labels has to be provided.'
        return self._attack_cls(self._task_def).attack(self._model, x, y)


class DefenseWrapper(TensorFlowV2Classifier):
    """ART wrapper for defenses."""

    def __init__(self, model):
        super().__init__(
            lambda x, **kwargs: model.classify(x),
            nb_classes=NUM_CLASSES,
            input_shape=INPUT_SHAPE,
            clip_values=(0., 1.),
            loss_object=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))

    def set_learning_phase(self, train):
        # Without overloading this method, armory does not work
        pass


def get_art_model(model_kwargs, wrapper_kwargs, weights_file):
    del wrapper_kwargs
    del weights_file
    defense_name = model_kwargs['defense']
    defense_path = get_defense_path(defense_name)

    model, _, _, dataset_name = load_defense_and_attack(