diff --git a/dreamify/classicdream.py b/dreamify/classicdream.py index b6872f1..817aeca 100644 --- a/dreamify/classicdream.py +++ b/dreamify/classicdream.py @@ -40,7 +40,7 @@ def classicdream( ): output_path = Path(output_path) - ft_ext = FeatureExtractor(model_name, dream_style="classic", layer_settings=None) + ft_ext = FeatureExtractor(model_name, "classic", layer_settings, channel_settings) get_tiled_gradients = TiledGradients(ft_ext.feature_extractor) img = get_image(image_path) diff --git a/dreamify/deepdream.py b/dreamify/deepdream.py index c72278a..3fdc42e 100644 --- a/dreamify/deepdream.py +++ b/dreamify/deepdream.py @@ -3,16 +3,12 @@ import IPython.display as display import tensorflow as tf -from dreamify.decorators import optional_global_determinism, validate_dream_params +from dreamify.decorators import (optional_global_determinism, + validate_dream_params) from dreamify.lib import Config, FeatureExtractor from dreamify.lib.models import TiledGradients -from dreamify.utils import ( - deprocess_image, - get_image, - preprocess_image, - save_output, - show, -) +from dreamify.utils import (deprocess_image, get_image, preprocess_image, + save_output, show) @optional_global_determinism @@ -37,7 +33,7 @@ def deepdream( ): output_path = Path(output_path) - ft_ext = FeatureExtractor(model_name, dream_style="deep", layer_settings=None) + ft_ext = FeatureExtractor(model_name, "deep", layer_settings, channel_settings) get_tiled_gradients = TiledGradients(ft_ext.feature_extractor) img = get_image(image_path) diff --git a/dreamify/dream.py b/dreamify/dream.py index 6e93f6d..010885c 100644 --- a/dreamify/dream.py +++ b/dreamify/dream.py @@ -45,9 +45,7 @@ def dream( ): output_path = Path(output_path) - ft_ext = FeatureExtractor( - model_name=model_name, dream_style="shallow", layer_settings=None - ) + ft_ext = FeatureExtractor(model_name, "shallow", layer_settings, channel_settings) original_img = get_image(image_path) original_img = preprocess_image(original_img, model_name) diff --git a/dreamify/lib/feature_extractor.py b/dreamify/lib/feature_extractor.py index b0d7790..c0f5aec 100644 --- a/dreamify/lib/feature_extractor.py +++ b/dreamify/lib/feature_extractor.py @@ -5,10 +5,11 @@ class FeatureExtractor: - def __init__(self, model_name, dream_style, layer_settings): + def __init__(self, model_name, dream_style, layer_settings, channel_settings): self.model, self.layer_settings = choose_base_model( model_name, dream_style, layer_settings ) + self.channel_settings = channel_settings # outputs is either: # - A list of layer outputs (for DeepDream) diff --git a/dreamify/lib/models/base/base_models.py b/dreamify/lib/models/base/base_models.py index 1928d8f..edea04d 100644 --- a/dreamify/lib/models/base/base_models.py +++ b/dreamify/lib/models/base/base_models.py @@ -1,10 +1,10 @@ import random from dreamify.lib.models.base.constants import MODEL_MAP, ModelType -from dreamify.lib.models.base.layer_settings import ModelLayerNames +from dreamify.lib.models.base.layer_names import ModelLayerNames -def get_layer_settings(model_name_enum: ModelType, dream_style="deep"): +def get_layer_names(model_name_enum: ModelType, dream_style): model_settings = ModelLayerNames[model_name_enum.name] if dream_style == "deep": return model_settings.deep @@ -27,7 +27,7 @@ def choose_base_model(model_name: str, dream_style="deep", layer_settings=None): model_fn = MODEL_MAP[model_name_enum] base_model = model_fn(weights="imagenet", include_top=False) - layer_names = get_layer_settings(model_name_enum, dream_style) + layer_names = get_layer_names(model_name_enum, dream_style) if layer_settings is not None: if layers_settings == "all": diff --git a/dreamify/lib/models/base/layer_names.py b/dreamify/lib/models/base/layer_names.py index 591f26b..7ebddca 100644 --- a/dreamify/lib/models/base/layer_names.py +++ b/dreamify/lib/models/base/layer_names.py @@ -10,18 +10,18 @@ class ModelLayerNames(Enum): INCEPTION_V3 = ( ["mixed3", "mixed5"], generate_shallow_settings(["mixed4", "mixed5", "mixed6", "mixed7"]), - [f"mixed_{i}" for i in range(1, 12)], + [f"mixed_{i}" for i in range(1, 12)] ) VGG19 = ( ["block5_conv3", "block5_conv2"], generate_shallow_settings( ["block5_conv3", "block5_conv2", "block4_conv3", "block3_conv3"] ), - [f"block1_conv{i}" for i in range(1, 3)] - + [f"block2_conv{i}" for i in range(1, 3)] - + [f"block3_conv{i}" for i in range(1, 6)] - + [f"block4_conv{i}" for i in range(1, 6)] - + [f"block5_conv{i}" for i in range(1, 6)], + [f"block1_conv{i}" for i in range(1, 3)] + + [f"block2_conv{i}" for i in range(1, 3)] + + [f"block3_conv{i}" for i in range(1, 6)] + + [f"block4_conv{i}" for i in range(1, 6)] + + [f"block5_conv{i}" for i in range(1, 6)] ) DENSENET121 = ( ["conv5_block16_1_conv", "conv4_block24_1_conv"], @@ -33,7 +33,7 @@ class ModelLayerNames(Enum): "conv2_block12_1_conv", ] ), - [], + [] ) EFFICIENTNET_V2L = ( ["block7a_project_bn", "block6a_expand_activation"], @@ -45,12 +45,12 @@ class ModelLayerNames(Enum): "block4a_expand_activation", ] ), - [], + [] ) INCEPTION_RESNET_V2 = ( ["mixed_7a", "mixed_6a"], generate_shallow_settings(["mixed_7a", "mixed_6a", "mixed_5a", "mixed_4a"]), - [], + [] ) RESNET152V2 = ( ["conv5_block3_out", "conv4_block6_out"], @@ -62,7 +62,7 @@ class ModelLayerNames(Enum): "conv2_block3_out", ] ), - [], + [] ) XCEPTION = ( ["block14_sepconv2_act", "block13_sepconv2_act"], @@ -74,7 +74,7 @@ class ModelLayerNames(Enum): "block11_sepconv2_act", ] ), - [], + [] ) CONVNEXT_XL = ( ["stage4_block2_depthwise_conv", "stage3_block2_depthwise_conv"], @@ -86,7 +86,7 @@ class ModelLayerNames(Enum): "stage1_block2_depthwise_conv", ] ), - [], + [] ) MOBILENET_V2 = ( ["block_16_depthwise", "block_13_depthwise"], @@ -98,7 +98,7 @@ class ModelLayerNames(Enum): "block_5_depthwise", ] ), - [], + [] ) @property @@ -115,3 +115,4 @@ def shallow(self): def classic(self): """Returns the convolutional blocks of the model""" return self.value[2] + \ No newline at end of file diff --git a/dreamify/lib/models/tiled_gradients.py b/dreamify/lib/models/tiled_gradients.py index 69ea591..01dc779 100644 --- a/dreamify/lib/models/tiled_gradients.py +++ b/dreamify/lib/models/tiled_gradients.py @@ -1,3 +1,5 @@ +from random import choice + import tensorflow as tf @@ -54,10 +56,26 @@ def calc_loss(img, feature_extractor): """Calculate the DeepDream loss by maximizing activations.""" img_batch = tf.expand_dims(img, axis=0) layer_activations = feature_extractor(img_batch) + + if feature_extractor.channel_settings == "random": + # Select a random channel for each layer + for layer_name in layer_activations: + channels = layer_activations[layer_name].shape[-1] + random_channel = choice(channels) + layer_activations[layer_name] = layer_activations[layer_name][ + ..., random_channel + ] + if len(layer_activations) == 1: layer_activations = [layer_activations] - return tf.reduce_sum([tf.math.reduce_mean(act) for act in layer_activations]) + # Ensure activations are iterable + if isinstance(layer_activations, dict): + activations_list = list(layer_activations.values()) + else: + activations_list = [layer_activations] + + return tf.reduce_sum([tf.reduce_mean(act) for act in activations_list]) @staticmethod def random_roll(img, maxroll): diff --git a/dreamify/utils/dream_utils/utils.py b/dreamify/utils/dream_utils/utils.py index cc28364..f843019 100644 --- a/dreamify/utils/dream_utils/utils.py +++ b/dreamify/utils/dream_utils/utils.py @@ -1,3 +1,5 @@ +from random import choice + import tensorflow as tf from tqdm import trange @@ -5,10 +7,20 @@ def compute_loss(input_image, config): features = config.feature_extractor(input_image) loss = tf.zeros(shape=()) + for name in features.keys(): coeff = config.layer_settings[name] activation = features[name] + + # Apply channel selection logic + if config.feature_extractor.channel_settings == "random": + channels = activation.shape[-1] + random_channel = choice(channels) + activation = activation[..., random_channel] + + # Compute loss on selected activations loss += coeff * tf.reduce_mean(tf.square(activation[:, 2:-2, 2:-2, :])) + return loss diff --git a/pyproject.toml b/pyproject.toml index 3bde6f9..d4a7385 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "dreamify" -version = "1.0.0" +version = "1.1.3" description = "" authors = [ {name = "kenny",email = "97374837+kankenny@users.noreply.github.com"} diff --git a/test b/test new file mode 100644 index 0000000..9daeafb --- /dev/null +++ b/test @@ -0,0 +1 @@ +test