# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Colour Perception in Humans and Machines
#
# In this notebook we compare the perceived similarity in colour for RetinaNet with that of Humans and the statistics of natural scenes. These plots correspond to Figure 5 from the paper.
#
# **Note**: Unfortunately we can't share all the models for this experiment (due to file size restrictions) so the first two cells can't be run in colab.
#
# ## Our Models

# +
''' Can't run in colab '''
import colorsys
import numpy as np
import torch
from training.model import BaselineModel 
from statistics.wavelength import wavelength_to_rgb

def response_to_wavelength(model, layer, wavelength):
    wavelength = torch.tensor([float(wavelength)], requires_grad=True)
    rgb = wavelength_to_rgb(wavelength, gamma=0.7)
    
    stimulus = torch.ones((1,32,32,3), requires_grad=False) * rgb
    stimulus = stimulus.permute(0, 3, 2, 1)
    
    resp = model.forward_to_layer(stimulus, layer).max()
    resp.backward()

    return wavelength.grad, rgb


def compute_response_to_wavelengths(model, layer, inc=0.5, threshold=1e-4):
    responses = []

    step = 10
    for i in np.concatenate((np.arange(395, 645, step), [644])):
        base = response_to_wavelength(model, layer, i)
        responses.append(base[0])

    return np.array(responses)


def get_distances(model_dir=None, pretrained=True, cmode='colour', d_vvs=[1, 2, 3, 4], n_bn=[1, 2, 4, 8, 16, 32]):
    distances = []
    
    for d in d_vvs:
        for n in n_bn:
            for t in range(10):
                n_ch = 1 if cmode == 'grey' else 3
                model = BaselineModel(n, d, n_ch)

                if pretrained:
                    try:
                        model.load_conv_dict(torch.load(f'{model_dir}/model_{n}_{d}_{t}.pt', map_location='cpu'))
                    except:
                        model.load_state_dict(torch.load(f'{model_dir}/model_{n}_{d}_{t}.pt', map_location='cpu'))
                        
                r = compute_response_to_wavelengths(model, f'retina_relu2')
                distances.append(r)

    distances = np.array(distances)
    return distances

# distances = get_distances('../models/colour')
rand_distances = get_distances(pretrained=False)
distort_distances = get_distances('../models/colour-distort')

distances_narrow = get_distances('../models/colour', n_bn=[1, 2, 4], d_vvs=[3, 4])
distances_wide = get_distances('../models/colour', n_bn=[8, 16, 32], d_vvs=[0, 1])

# +
from scipy import stats
# %matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import rc
import matplotlib.font_manager
rc('font',**{'family':'serif','serif':['Computer Modern Roman'],'size':13})
rc('text', usetex=True)

plt.figure(figsize=(3,2.5))

sems = stats.sem(np.abs(distances_narrow), axis=0)
means = (np.abs(distances_narrow)).mean(axis=0)

w = np.concatenate((np.arange(395, 645, 10), [644]))
plt.plot(w, means, linestyle='-')
plt.fill_between(w, means - sems, means + sems, alpha=0.2)

plt.xlim(400, 700)
plt.gca().get_yaxis().set_ticks([])

plt.xlabel('Wavelength')
plt.ylabel('Sensitivity')
plt.savefig('figures/similarity_narrow.pdf', bbox_inches='tight')

# +
from scipy import stats
# %matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import rc
import matplotlib.font_manager
rc('font',**{'family':'serif','serif':['Computer Modern Roman'],'size':13})
rc('text', usetex=True)

plt.figure(figsize=(3,2.5))

sems = stats.sem(np.abs(distances_wide), axis=0)
means = (np.abs(distances_wide)).mean(axis=0)

w = np.concatenate((np.arange(395, 645, 10), [644]))
plt.plot(w, means, linestyle='-')
plt.fill_between(w, means - sems, means + sems, alpha=0.2)

plt.xlim(400, 700)
plt.gca().get_yaxis().set_ticks([])

plt.xlabel('Wavelength')
plt.ylabel('Sensitivity')
plt.savefig('figures/similarity_wide.pdf', bbox_inches='tight')

# +
''' Can't run in colab '''
plt.figure(figsize=(3,2.5))

sems = stats.sem(np.abs(rand_distances), axis=0)
means = (np.abs(rand_distances)).mean(axis=0)

w = np.concatenate((np.arange(395, 645, 10), [644]))
plt.plot(w, means, linestyle='-', color='C1')
plt.fill_between(w, means - sems, means + sems, alpha=0.2, facecolor='C1')

plt.xlim(400, 700)
plt.gca().get_yaxis().set_ticks([])

plt.xlabel('Wavelength')
plt.ylabel('Sensitivity')
plt.savefig('figures/similarity_random.pdf', bbox_inches='tight')

# +
plt.figure(figsize=(3,2.5))

sems = stats.sem(np.abs(distort_distances), axis=0)
means = (np.abs(distort_distances)).mean(axis=0)

w = np.concatenate((np.arange(395, 645, 10), [644]))
plt.plot(w, means, linestyle='-')
plt.fill_between(w, means - sems, means + sems, alpha=0.2)

plt.xlim(400, 700)
plt.gca().get_yaxis().set_ticks([])

plt.xlabel('Wavelength')
plt.ylabel('Sensitivity')
plt.savefig('figures/similarity_distort.pdf', bbox_inches='tight')

# +
# %matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import rc
import matplotlib.font_manager
rc('font',**{'family':'serif','serif':['Computer Modern Roman'],'size':13})
rc('text', usetex=True)
from matplotlib.colors import ListedColormap
from matplotlib import cm
from statistics.wavelength import wavelength_to_rgb
import numpy as np

rs = []
gs = []
bs = []

ws = list(range(400, 701))

for lam in ws:
    rgb = wavelength_to_rgb(lam)
    
    rs.append(rgb[0])
    gs.append(rgb[1])
    bs.append(rgb[2])

plt.figure(figsize=(9,2.8))
plt.box(False)
plt.plot(ws, rs, linestyle='-.', color='r')
plt.plot(ws, gs, linestyle='--', color='g')
plt.plot(ws, bs, linestyle=':', color='b')

plt.legend(['Red', 'Green', 'Blue'], frameon=False)

colours = np.stack((np.array(rs), np.array(gs), np.array(bs)), axis=1)

colours = ListedColormap(colours)
cb = plt.gcf().colorbar(cm.ScalarMappable(cmap=colours), pad=0.27, orientation='horizontal', ticks=[], aspect=25)
cb.outline.set_visible(False)

plt.yticks([])
plt.xlim(400, 700)
plt.ylim(-0.05, 1.05)
plt.xlabel('Wavelength')
plt.savefig('figures/wavelength.pdf', bbox_inches='tight')

# +
# %matplotlib inline
import matplotlib.pyplot as plt

from colour.plotting import plot_RGB_colourspaces_in_chromaticity_diagram_CIE1931

fig, ax = plot_RGB_colourspaces_in_chromaticity_diagram_CIE1931(['CIE-LAB'], legend=False, standalone=False, axes_visible=False, title='', spectral_locus_colours='RGB', diagram_opacity=0.7)
plt.legend(['RGB'], frameon=False)
# -

# ## Load Dependencies - Colab Only

from os.path import exists
if not exists('opponency.zip'):
    # !wget -O opponency.zip https://github.com/ecs-vlc/opponency/archive/master.zip
    # !unzip -qq opponency.zip
    # !mv opponency-master/* ./
    # !rm -r opponency-master

# ## Other Plots

# +
# %matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import rc
import matplotlib.font_manager
rc('font',**{'family':'serif','serif':['Computer Modern Roman'],'size':13})
rc('text', usetex=True)
import pandas as pd
bedford = pd.read_csv('bedford1958.csv', names=['wavelength', 'sensitivity'], header=None)
bedford[bedford.wavelength.between(380, 650)].plot('wavelength', 'sensitivity', legend=False, figsize=(3,2.5), linestyle='-')

plt.xlim(400, 700)
plt.gca().get_yaxis().set_ticks([])
plt.xlabel('Wavelength')
plt.ylabel('Perceptual Similarity')
plt.savefig('figures/similarity_bedford.pdf', bbox_inches='tight')

# +
# %matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import rc
import matplotlib.font_manager
rc('font',**{'family':'serif','serif':['Computer Modern Roman'],'size':13})
rc('text', usetex=True)
import pandas as pd

long = pd.read_csv('long2006.csv', names=['wavelength', 'sensitivity'], header=None)
# long = long[long.wavelength.between(420, 650)]  # .plot('wavelength', 'sensitivity')
# bedford
import numpy as np
p = np.poly1d(np.polyfit(long.wavelength, long.sensitivity, 8))
import matplotlib.pyplot as plt

x = np.linspace(430, 600, 1000)

plt.figure(figsize=(3,2.5))

plt.plot(x, p(x), linestyle='-')
plt.xlim(400, 700)
plt.gca().get_yaxis().set_ticks([])
plt.xlabel('Wavelength')
plt.ylabel('Predicted Similarity')
plt.savefig('figures/similarity_long.pdf', bbox_inches='tight')
# -



StarcoderdataJupyterScriptsDedupFiltered