<commit_before>"""Deep Dreaming using Caffe and Google's Inception convolutional neural network."""

# pylint: disable=invalid-name, wrong-import-position

from collections import namedtuple, OrderedDict
import os
from pathlib import Path
import re

os.environ['GLOG_minloglevel'] = '1'
import caffe
import numpy as np
from PIL import Image
from scipy import ndimage
from tqdm import tqdm

EPS = np.finfo(np.float32).eps

CNNData = namedtuple('CNNData', 'deploy model mean categories')
CNNData.__new__.__defaults__ = (None,)  # Make categories optional.

_BASE_DIR = Path(__file__).parent.parent
GOOGLENET_BVLC = CNNData(
    _BASE_DIR/'bvlc_googlenet/deploy.prototxt',
    _BASE_DIR/'bvlc_googlenet/bvlc_googlenet.caffemodel',
    (104, 117, 123),
    categories=_BASE_DIR/'bvlc_googlenet/categories.txt')
GOOGLENET_PLACES205 = CNNData(
    _BASE_DIR/'googlenet_places205/deploy_places205.prototxt',
    _BASE_DIR/'googlenet_places205/googlelet_places205_train_iter_2400000.caffemodel',
    (104.051, 112.514, 116.676),  # TODO: find the actual Places205 mean
    categories=_BASE_DIR/'googlenet_places205/categories.txt')
GOOGLENET_PLACES365 = CNNData(
    _BASE_DIR/'googlenet_places365/deploy_googlenet_places365.prototxt',
    _BASE_DIR/'googlenet_places365/googlenet_places365.caffemodel',
    (104.051, 112.514, 116.676),
    categories=_BASE_DIR/'googlenet_places365/categories_places365.txt')


def to_image(arr):
    """Clips the values in a float32 ndarray to 0-255 and converts it to a PIL image."""
    return Image.fromarray(np.uint8(np.clip(np.round(arr), 0, 255)))


def _resize(arr, size, method=Image.BICUBIC):
    h, w = size
    arr = np.float32(arr)
    if arr.ndim == 3:
        planes = [arr[i, :, :] for i in range(arr.shape[0])]
    else:
        raise TypeError('Only 3D CxHxW arrays are supported')
    imgs = [Image.fromarray(plane) for plane in planes]
    imgs_resized = [img.resize((w, h), method) for img in imgs]
    return np.stack([np.array(img) for img in imgs_resized])


class ShapeError(Exception):
    """Raised by CNN when an invalid layer shape is requested which would otherwise crash Caffe."""
    def __str__(self):
        return 'bad shape %s' % self.args


class _LayerIndexer:
    def __init__(self, net, attr):
        self.net, self.attr = net, attr

    def __getitem__(self, key):
        return getattr(self.net.blobs[key], self.attr)[0]

    def __setitem__(self, key, value):
        getattr(self.net.blobs[key], self.attr)[0] = value


class CNN:
    """Represents an instance of a Caffe convolutional neural network."""

    def __init__(self, cnndata, gpu=None):
        """Initializes a CNN.

        Args:
            gpu (Optional[int]): If present, Caffe will use this GPU device number. On a typical
                system with one GPU, it should be 0. If not present Caffe will use the CPU.
        """
        self.net = caffe.Classifier(str(cnndata.deploy), str(cnndata.model),
                                    mean=np.float32(cnndata.mean), channel_swap=(2, 1, 0))
        self.data = _LayerIndexer(self.net, 'data')
        self.diff = _LayerIndexer(self.net, 'diff')
        self.categories = [str(i) for i in range(self.data['prob'].size)]
        if cnndata.categories is not None:
            self.categories = open(str(cnndata.categories)).read().splitlines()
        self.img = np.zeros_like(self.data['data'])
        self.total_px = 0
        self.progress_bar = None
        if gpu is not None:
            caffe.set_device(gpu)
            caffe.set_mode_gpu()
        else:
            caffe.set_mode_cpu()

    def _preprocess(self, img):
        return np.rollaxis(np.float32(img), 2)[::-1] - self.net.transformer.mean['data']

    def _deprocess(self, img):
        return np.dstack((img + self.net.transformer.mean['data'])[::-1])

    def get_features(self, input_img, layers, max_tile_size=512):
        """Retrieve feature maps from the classification (forward) phase of operation.

        Example:
            cnn.get_features(img, ['prob'])['prob'] classifies 'img' and returns the predicted
            probability distribution over the network's categories.

        Returns:
            A dict which maps each layer in layers to a retrieved feature map.
        """
        input_arr = self._preprocess(np.float32(input_img))
        h = min(max_tile_size, input_arr.shape[1])
        w = min(max_tile_size, input_arr.shape[2])
        if max(*input_arr.shape[1:]) > max_tile_size:
            input_arr = _resize(input_arr, (h, w))
        self.net.blobs['data'].reshape(1, 3, h, w)
        self.data['data'] = input_arr
        end = self.layers()[-1]
        self.net.forward(end=end)
        return {layer: self.data[layer].copy() for layer in layers}

    def _grad_tiled(self, layers, progress=False, max_tile_size=512, auto_weight=True):
        # pylint: disable=too-many-locals
        if progress:
            if not self.progress_bar:
                self.progress_bar = tqdm(
                    total=self.total_px, unit='pix', unit_scale=True, ncols=80, smoothing=0.1)

        h, w = self.img.shape[1:]  # Height and width of input image
        ny, nx = (h-1)//max_tile_size+1, (w-1)//max_tile_size+1  # Number of tiles per dimension
        g = np.zeros_like(self.img)
        for y in range(ny):
            th = h//ny
            if y == ny-1:
                th += h - th*ny
            for x in range(nx):
                tw = w//nx
                if x == nx-1:
                    tw += w - tw*nx
                self.net.blobs['data'].reshape(1, 3, th, tw)
                sy, sx = h//ny*y, w//nx*x
                self.data['data'] = self.img[:, sy:sy+th, sx:sx+tw]

                for layer in layers.keys():
                    self.diff[layer] = 0
                self.net.forward(end=next(iter(layers.keys())))
                layers_list = list(layers.keys())
                for i, layer in enumerate(layers_list):
                    if auto_weight:
                        self.diff[layer] += \
                            self.data[layer]*layers[layer]/np.abs(self.data[layer]).sum()
                    else:
                        self.diff[layer] += self.data[layer]*layers[layer]
                    if i+1 == len(layers):
                        self.net.backward(start=layer)
                    else:
                        self.net.backward(start=layer, end=layers_list[i+1])

                g[:, sy:sy+th, sx:sx+tw] = self.diff['data']

                if progress:
                    self.progress_bar.update(th*tw)
        return g

    def _step(self, n=1, step_size=1.5, jitter=32, seed=0, smooth=0, **kwargs):
        np.random.seed(self.img.size + seed)
        for _ in range(n):
            x, y = np.random.randint(-jitter, jitter+1, 2)
            self.img = np.roll(np.roll(self.img, x, 2), y, 1)
            g = self._grad_tiled(**kwargs)
            g /= np.mean(np.abs(g)) + EPS
            if smooth:
                kernel = np.float32([[[0, 1, 0], [1, -4, 1], [0, 1, 0]]])
                g2 = ndimage.convolve(self.img, kernel)
                g2 /= np.mean(np.abs(g2)) + EPS
                g += g2 * smooth
                self.img += step_size * g / (np.mean(np.abs(g)) + EPS)
            else:
                self.img += step_size * g
            self.img = np.roll(np.roll(self.img, -x, 2), -y, 1)

    def _octave_detail(self, base, scales=4, min_size=32, n=10, per_octave=2, **kwargs):
        if min(base.shape[1:]) < 32:
            raise ShapeError(base.shape)
        factor = 2**(1/per_octave)
        detail = np.zeros_like(base, dtype=np.float32)
        self.total_px += base.shape[1] * base.shape[2] * n
        if scales != 1:
            hf, wf = np.int32(np.round(np.array(base.shape)[-2:]/factor))
            if min(hf, wf) >= min_size:
                smaller_base = _resize(base, (hf, wf))
                smaller_detail = self._octave_detail(
                    smaller_base, scales-1, min_size, n, per_octave, **kwargs)
                detail = _resize(smaller_detail, base.shape[-2:])
        self.img = base + detail
        self._step(n, **kwargs)
        return self.img - base

    def layers(self, pattern='.*'):<commit_msg>Minor tweak to smoothing<commit_after>"""Deep Dreaming using Caffe and Google's Inception convolutional neural network."""

# pylint: disable=invalid-name, wrong-import-position

from collections import namedtuple, OrderedDict
import os
from pathlib import Path
import re

os.environ['GLOG_minloglevel'] = '1'
import caffe
import numpy as np
from PIL import Image
from scipy import ndimage
from tqdm import tqdm

EPS = np.finfo(np.float32).eps

CNNData = namedtuple('CNNData', 'deploy model mean categories')
CNNData.__new__.__defaults__ = (None,)  # Make categories optional.

_BASE_DIR = Path(__file__).parent.parent
GOOGLENET_BVLC = CNNData(
    _BASE_DIR/'bvlc_googlenet/deploy.prototxt',
    _BASE_DIR/'bvlc_googlenet/bvlc_googlenet.caffemodel',
    (104, 117, 123),
    categories=_BASE_DIR/'bvlc_googlenet/categories.txt')
GOOGLENET_PLACES205 = CNNData(
    _BASE_DIR/'googlenet_places205/deploy_places205.prototxt',
    _BASE_DIR/'googlenet_places205/googlelet_places205_train_iter_2400000.caffemodel',
    (104.051, 112.514, 116.676),  # TODO: find the actual Places205 mean
    categories=_BASE_DIR/'googlenet_places205/categories.txt')
GOOGLENET_PLACES365 = CNNData(
    _BASE_DIR/'googlenet_places365/deploy_googlenet_places365.prototxt',
    _BASE_DIR/'googlenet_places365/googlenet_places365.caffemodel',
    (104.051, 112.514, 116.676),
    categories=_BASE_DIR/'googlenet_places365/categories_places365.txt')


def to_image(arr):
    """Clips the values in a float32 ndarray to 0-255 and converts it to a PIL image."""
    return Image.fromarray(np.uint8(np.clip(np.round(arr), 0, 255)))


def _resize(arr, size, method=Image.BICUBIC):
    h, w = size
    arr = np.float32(arr)
    if arr.ndim == 3:
        planes = [arr[i, :, :] for i in range(arr.shape[0])]
    else:
        raise TypeError('Only 3D CxHxW arrays are supported')
    imgs = [Image.fromarray(plane) for plane in planes]
    imgs_resized = [img.resize((w, h), method) for img in imgs]
    return np.stack([np.array(img) for img in imgs_resized])


class ShapeError(Exception):
    """Raised by CNN when an invalid layer shape is requested which would otherwise crash Caffe."""
    def __str__(self):
        return 'bad shape %s' % self.args


class _LayerIndexer:
    def __init__(self, net, attr):
        self.net, self.attr = net, attr

    def __getitem__(self, key):
        return getattr(self.net.blobs[key], self.attr)[0]

    def __setitem__(self, key, value):
        getattr(self.net.blobs[key], self.attr)[0] = value


class CNN:
    """Represents an instance of a Caffe convolutional neural network."""

    def __init__(self, cnndata, gpu=None):
        """Initializes a CNN.

        Args:
            gpu (Optional[int]): If present, Caffe will use this GPU device number. On a typical
                system with one GPU, it should be 0. If not present Caffe will use the CPU.
        """
        self.net = caffe.Classifier(str(cnndata.deploy), str(cnndata.model),
                                    mean=np.float32(cnndata.mean), channel_swap=(2, 1, 0))
        self.data = _LayerIndexer(self.net, 'data')
        self.diff = _LayerIndexer(self.net, 'diff')
        self.categories = [str(i) for i in range(self.data['prob'].size)]
        if cnndata.categories is not None:
            self.categories = open(str(cnndata.categories)).read().splitlines()
        self.img = np.zeros_like(self.data['data'])
        self.total_px = 0
        self.progress_bar = None
        if gpu is not None:
            caffe.set_device(gpu)
            caffe.set_mode_gpu()
        else:
            caffe.set_mode_cpu()

    def _preprocess(self, img):
        return np.rollaxis(np.float32(img), 2)[::-1] - self.net.transformer.mean['data']

    def _deprocess(self, img):
        return np.dstack((img + self.net.transformer.mean['data'])[::-1])

    def get_features(self, input_img, layers, max_tile_size=512):
        """Retrieve feature maps from the classification (forward) phase of operation.

        Example:
            cnn.get_features(img, ['prob'])['prob'] classifies 'img' and returns the predicted
            probability distribution over the network's categories.

        Returns:
            A dict which maps each layer in layers to a retrieved feature map.
        """
        input_arr = self._preprocess(np.float32(input_img))
        h = min(max_tile_size, input_arr.shape[1])
        w = min(max_tile_size, input_arr.shape[2])
        if max(*input_arr.shape[1:]) > max_tile_size:
            input_arr = _resize(input_arr, (h, w))
        self.net.blobs['data'].reshape(1, 3, h, w)
        self.data['data'] = input_arr
        end = self.layers()[-1]
        self.net.forward(end=end)
        return {layer: self.data[layer].copy() for layer in layers}

    def _grad_tiled(self, layers, progress=False, max_tile_size=512, auto_weight=True):
        # pylint: disable=too-many-locals
        if progress:
            if not self.progress_bar:
                self.progress_bar = tqdm(
                    total=self.total_px, unit='pix', unit_scale=True, ncols=80, smoothing=0.1)

        h, w = self.img.shape[1:]  # Height and width of input image
        ny, nx = (h-1)//max_tile_size+1, (w-1)//max_tile_size+1  # Number of tiles per dimension
        g = np.zeros_like(self.img)
        for y in range(ny):
            th = h//ny
            if y == ny-1:
                th += h - th*ny
            for x in range(nx):
                tw = w//nx
                if x == nx-1:
                    tw += w - tw*nx
                self.net.blobs['data'].reshape(1, 3, th, tw)
                sy, sx = h//ny*y, w//nx*x
                self.data['data'] = self.img[:, sy:sy+th, sx:sx+tw]

                for layer in layers.keys():
                    self.diff[layer] = 0
                self.net.forward(end=next(iter(layers.keys())))
                layers_list = list(layers.keys())
                for i, layer in enumerate(layers_list):
                    if auto_weight:
                        self.diff[layer] += \
                            self.data[layer]*layers[layer]/np.abs(self.data[layer]).sum()
                    else:
                        self.diff[layer] += self.data[layer]*layers[layer]
                    if i+1 == len(layers):
                        self.net.backward(start=layer)
                    else:
                        self.net.backward(start=layer, end=layers_list[i+1])

                g[:, sy:sy+th, sx:sx+tw] = self.diff['data']

                if progress:
                    self.progress_bar.update(th*tw)
        return g

    def _step(self, n=1, step_size=1.5, jitter=32, seed=0, smooth=0, **kwargs):
        np.random.seed(self.img.size + seed)
        for _ in range(n):
            x, y = np.random.randint(-jitter, jitter+1, 2)
            self.img = np.roll(np.roll(self.img, x, 2), y, 1)
            g = self._grad_tiled(**kwargs)
            g /= np.mean(np.abs(g)) + EPS
            if smooth:
                kernel = np.float32([[[0, -1, 0], [-1, 4, -1], [0, -1, 0]]])
                g2 = ndimage.convolve(self.img, kernel)
                g2 /= np.mean(np.abs(g2)) + EPS
                g -= g2 * smooth
                self.img += step_size * g / (np.mean(np.abs(g)) + EPS)
            else:
                self.img += step_size * g
            self.img = np.roll(np.roll(self.img, -x, 2), -y, 1)

    def _octave_detail(self, base, scales=4, min_size=32, n=10, per_octave=2, **kwargs):
        if min(base.shape[1:]) < 32:
            raise ShapeError(base.shape)
        factor = 2**(1/per_octave)
        detail = np.zeros_like(base, dtype=np.float32)
        self.total_px += base.shape[1] * base.shape[2] * n
        if scales != 1:
            hf, wf = np.int32(np.round(np.array(base.shape)[-2:]/factor))
            if min(hf, wf) >= min_size:
                smaller_base = _resize(base, (hf, wf))
                smaller_detail = self._octave_detail(
                    smaller_base, scales-1, min_size, n, per_octave, **kwargs)
                detail = _resize(smaller_detail, base.shape[-2:])
        self.img = base + detail
        self._step(n, **kwargs)
        return self.img - base

    def layers(self, pattern='.*'):
{'fraction_non_alphanumeric': 0.1301361536591296, 'fraction_numerical': 0.025042548018477998, 'mean_word_length': 2.2438880126182967, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3480918', 'n_tokens_mistral': 5816, 'n_tokens_neox': 5348, 'n_words': 1242}