Merge pull request #1 from KordingLab/feature/qc-single-patch-orientation

Thank you!

Author: aribenjamin

single_patch_orientation/README.md

@@ -0,0 +1,19 @@
+# Single patch orientation using pre-trained AlexNet
+
+Test pre-trained AlexNet with various orientation stimuli used in psychophysics to see if the model shows similar behaviors as in human observers.
+
+## Usage
+The main Python file for this project is `alexnet_to_orientation.py`.
+An example usage of this script is:
+
+```
+ $ python alexnet_to_orientation.py --epochs 10 \
+   --save-model --model-name 'alexNet_broadband_multiorinoise_naturaloriprior' \
+   --if-more-noise-levels
+```
+
+For all options and command-line arguments, please use:
+
+```
+ $ python alexnet_to_orientation.py -h
+```

single_patch_orientation/alexnet_to_orientation.py

@@ -0,0 +1,131 @@
+import numpy as np
+
+import matplotlib.pyplot as plt
+import argparse
+
+import torch
+import torch.optim as optim
+
+import orientation_stim
+import toy_model_train_orientation
+
+
+"""
+example of use:
+    python alexnet_to_orientation.py --epochs 10 \
+    --save-model --model-name 'alexNet_broadband_multiorinoise_naturaloriprior' \
+    --if-more-noise-levels 
+"""
+
+def main():
+    parser = argparse.ArgumentParser(description='Orientation with pre-trained alexnet')
+    parser.add_argument('--input-img-size', type=int, default=224, metavar='N',
+                        help='input image size (default: 224)')
+    parser.add_argument('--batch-size', type=int, default=128, metavar='N',
+                        help='input batch size for training (default: 128)')
+    parser.add_argument('--epoch-size', type=int, default=100, metavar='N',
+                        help='epoch size (# of batches) for generating images online.')
+    parser.add_argument('--test-epoch-size', type=int, default=20, metavar='N',
+                        help='epoch size (# of batches) for generating images online.')
+    parser.add_argument('--epochs', type=int, default=10, metavar='N',
+                        help='number of epochs to train (default: 10)')
+    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
+                        help='how many batches to wait before logging training status')
+    parser.add_argument('--no-cuda', action='store_true', default=False,
+                        help='disables CUDA training')
+    parser.add_argument('--seed', type=int, default=1, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--save-model', action='store_true', default=False,
+                        help='For Saving the current Model')
+    parser.add_argument('--model-name', type=str, default='orient_cnn',
+                        help='Name of the current Model for saving')
+    parser.add_argument('--if-unif', action='store_true', default=False,  # use natural prior by default
+                        help='if the training distribution uniform')
+    parser.add_argument('--if-more-noise-levels', action='store_true', default=False,
+                        help='if multiple orientation noise levels used in training')
+
+    args = parser.parse_args()
+    args.vis_name = args.model_name
+
+    use_cuda = not args.no_cuda and torch.cuda.is_available()
+
+    torch.manual_seed(args.seed)
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    model = toy_model_train_orientation.SlimAlexNet(max_pool_layer_index=1,
+                                          last_layer_num_params=2).to(device)
+
+    # the parameters would be finetuned - should only be conv_pool.5
+    params_to_update = model.parameters()
+    print("Params to learn:")
+    for name, param in model.named_parameters():
+        if param.requires_grad == True:
+            print("\t", name)
+
+    optimizer_ft = optim.Adam(params_to_update, lr=1e-4)
+
+    test_loss_history = []
+    report_epoch = [1, 2, 4, 8, 16]
+    for epoch in range(1, args.epochs + 1):
+        toy_model_train_orientation.train(args, model, device,
+                                          optimizer_ft, epoch, if_alexNet=True)
+        test_loss = toy_model_train_orientation.test(args, model, device,
+                                                     if_alexNet=True)
+        test_loss_history.append(test_loss)
+        # save intermediate result figures
+        if epoch in report_epoch:
+            if (args.save_model):
+                torch.save(model.state_dict(),
+                           args.model_name + '_epoch' + str(epoch) + '.pt')
+                # save a figure with bias
+                model_file = args.model_name + '_epoch' + str(epoch) + '.pt'
+                ave_ori, all_ave_bias = toy_model_train_orientation.compare_bias(model_file,
+                                                                                 img_size=224, if_alexNet=True)
+                plt.figure(figsize=(5, 3))
+                plt.plot(ave_ori, all_ave_bias[:, 1] - all_ave_bias[:, 0])
+                plt.plot(ave_ori, all_ave_bias[:, 2] - all_ave_bias[:, 0])
+                plt.plot([np.pi / 4, np.pi / 4], [-5 / 180 * np.pi, 5 / 180 * np.pi], 'k--')
+                plt.plot([np.pi / 2, np.pi / 2], [-5 / 180 * np.pi, 5 / 180 * np.pi], 'k--')
+                plt.plot([3 * np.pi / 4, 3 * np.pi / 4], [-5 / 180 * np.pi, 5 / 180 * np.pi], 'k--')
+                plt.plot([0, np.pi], [0, 0], 'k--')
+                plt.xlim([0, np.pi])
+                plt.ylim(-5 / 180 * np.pi, 5 / 180 * np.pi)
+                plt.savefig(args.vis_name + '_epoch' + str(epoch) + '.pdf')
+
+    # save final report if not has been saved
+    if args.epochs not in report_epoch:
+        if (args.save_model):
+            torch.save(model.state_dict(),
+                       args.model_name + '_epoch' + str(epoch) + '.pt')
+
+    # save a final figure with cosine similarity and bias
+    feature_cosSim = toy_model_train_orientation.feature_similarity(args.input_img_size,
+                                                   model, feature_layer_ind=3, if_alexNet=True)
+    plt.figure(figsize=(8, 3))
+    plt.subplot(1, 2, 1)
+    plt.plot(np.arange(179), feature_cosSim)
+    plt.xlim([0, 180])
+    # bias differences
+    model_file = args.model_name + '_epoch' + str(epoch) + '.pt'
+    ave_ori, all_ave_bias = toy_model_train_orientation.compare_bias(model_file,
+                                                                     img_size=224, if_alexNet=True)
+    plt.subplot(1, 2, 2)
+    plt.plot(ave_ori, all_ave_bias[:, 1] - all_ave_bias[:, 0])
+    plt.plot(ave_ori, all_ave_bias[:, 2] - all_ave_bias[:, 0])
+    plt.plot([0, np.pi], [0, 0], 'k--')
+    plt.plot([np.pi / 4, np.pi / 4], [-5 / 180 * np.pi, 5 / 180 * np.pi], 'k--')
+    plt.plot([np.pi / 2, np.pi / 2], [-5 / 180 * np.pi, 5 / 180 * np.pi], 'k--')
+    plt.plot([3 * np.pi / 4, 3 * np.pi / 4], [-5 / 180 * np.pi, 5 / 180 * np.pi], 'k--')
+    plt.xlim([0, np.pi])
+    plt.ylim(-5 / 180 * np.pi, 5 / 180 * np.pi)
+    plt.savefig(args.vis_name + '_final_summary' + '.pdf')
+
+    # figure showing the test loss progress
+    plt.figure()
+    plt.plot(np.arange(1, args.epochs + 1), test_loss_history)
+    plt.savefig(args.vis_name + '_loss_history' + '.pdf')
+
+
+if __name__ == '__main__':
+    main()

single_patch_orientation/orientation_stim.py

@@ -0,0 +1,196 @@
+import numpy as np
+import scipy.ndimage
+import matplotlib
+
+# matplotlib.use("TkAgg")
+import matplotlib.pyplot as plt
+
+# torch package
+import torch
+from torch.distributions import normal
+
+
+def grating(size=500, pixelsPerDegree=200, spatial_freq=3, spatial_phase=0,
+            orientation=np.pi/4, contrast=1):
+    '''
+    The output range is -1 to 1 if contrast is 1
+
+    size: number of pixel of the image patch, assuming square
+    spatial_freq: cycle per visual angle
+    spatial_phase: in radians
+    pixelsPerDegree: number of patch pixels in one degree visual angle
+    orientation: in radians
+    '''
+    x, y = np.meshgrid(np.arange(size), np.arange(size))
+    x = (x - size / 2.0) / pixelsPerDegree
+    y = (y - size / 2.0) / pixelsPerDegree
+
+    return contrast * np.cos(spatial_phase +
+                      2 * np.pi * spatial_freq *
+                      (x * np.sin(orientation) + y * np.cos(orientation)))
+
+
+def gabor(size=500, pixelsPerDegree=100, spatial_freq=3, spatial_phase=0,
+          orientation=np.pi/4, contrast=1, sigma=.5, spatial_aspect_ratio=1):
+    '''
+    adds an exponential modulation (Gaussian envelope) on top of "grating"
+
+    spatial_freq: cycle/visual angle
+    spatial_phase: in radians
+    orientation: in radians
+    contrast: if 1, values range from -1 to 1
+    sigma: standard deviation of the Gaussian envelope (in visual angle)
+    spatial_aspect_ratio (gamma): if not 1, distorted along orientation.
+           specifies the ellipticity of the support of the Gabor function.
+
+    '''
+
+    x, y = np.meshgrid(np.arange(size), np.arange(size))
+    x = (x - size / 2.0) / pixelsPerDegree #in visual angle
+    y = (y - size / 2.0) / pixelsPerDegree
+
+    # rotation
+    x_theta = x * np.cos(orientation) + y * np.sin(orientation)
+    y_theta = -x * np.sin(orientation) + y * np.cos(orientation)
+
+    # Gaussian envelope
+    sigma_x = sigma
+    sigma_y = float(sigma) / spatial_aspect_ratio
+    gaussian_envelope = np.exp(-.5 * (x_theta ** 2 / sigma_x ** 2 +
+                                      y_theta ** 2 / sigma_y ** 2))
+
+    # sinusoidal grating
+    grating = np.cos(2 * np.pi * spatial_freq * x_theta + spatial_phase)
+    gabor = gaussian_envelope * grating
+
+    # normalize for contrast
+    gabor_min = np.min(gabor[:])
+    gabor_max = np.max(gabor[:])
+    gabor = (gabor - gabor_min) * contrast * 2 / (gabor_max - gabor_min) - 1
+
+    return gabor
+
+
+def matlab_style_gauss2D(shape=(3, 3), sigma=0.5):
+    """
+    2D gaussian mask - should give the same result as MATLAB's
+    fspecial('gaussian',[shape],[sigma])
+    """
+    m,n = [(ss-1.)/2. for ss in shape]
+    y,x = np.ogrid[-m:m+1,-n:n+1]
+    h = np.exp( -(x*x + y*y) / (2.*sigma*sigma))
+    h[ h < np.finfo(h.dtype).eps*h.max() ] = 0
+    sumh = h.sum()
+    if sumh != 0:
+        h /= sumh
+    return h
+
+
+def circular_mask(size=500, pixelsPerDegree=200, radius=1, polarity_in=1, polarity_out=0,
+                  if_filtered=False, filter_size = (15, 15), filter_width = 2):
+    '''
+    :param size: size of the image patch
+    :param radius: in visual angle
+    :param polarity_in: 1 or 0 inside the circle
+    :param filter_size and filter_width are in pixel units
+    :return: the mask
+    '''
+    x, y = np.meshgrid(np.arange(size), np.arange(size))
+    x = (x - size / 2.0) / pixelsPerDegree
+    y = (y - size / 2.0) / pixelsPerDegree
+
+    mask = np.ones([size, size]) * polarity_out
+    mask[np.sqrt(np.power(x, 2) + np.power(y, 2)) < radius] = polarity_in
+
+    # Gaussian filtering the mask
+    if if_filtered:
+        H = matlab_style_gauss2D(filter_size, filter_width)  # lowpass filter
+        mask = scipy.ndimage.convolve(mask, H, mode='nearest')  # filter
+
+    return mask
+
+# --- example grating ---
+# img = grating(500, 200, 3, np.pi/4, 1, np.pi/5)
+# mask = circular_mask(500, 200, 1,
+#                      if_filtered=True, filter_size=(50, 50), filter_width=10)
+# img = np.multiply(img, mask)
+# plt.imshow(img, cmap=plt.gray()), plt.show()
+
+
+def broadband_noise(size=64, contrast=1, if_low_pass=True, center_sf=0, sf_sigma=10,
+                    if_band_pass=False, low_sf=1.67, high_sf=10.67,
+                    orientation=20/180*np.pi, orient_sigma=10/180*np.pi):
+    """
+    broadband noise with either low pass [--low-pass] or
+                                band pass [--band-pass] spatial frequency;
+                                if set one as true, need to make sure the other is false
+
+    Args:
+        size: if size is not power of 2, crop to be the intended size
+        contrast:
+        center_sf: if 0, low pass
+        sf_sigma: in pixel; if Inf, all sf included (but exclude the corner)
+        low_sf: cycle/image
+        high_sf: cycle/image
+        orientation: in radians, note that the available range of orientation is 0 - pi
+        orient_sigma: in radians; if Inf, all orientations are included
+
+    Return:
+    """
+
+    size_tmp = np.power(2, np.ceil(np.log2(size))).astype(int)  # then crop to intended size
+    input_img = np.random.uniform(0, 1, [size_tmp, size_tmp])
+    max_sf = size_tmp / 2
+    img_center = np.matlib.repmat(np.floor(size_tmp / 2), 1, 2)
+
+    # Fourier transform and separate magnitude and phase
+    f = np.fft.fftshift(np.fft.fft2(input_img))
+    # mag_f = np.abs(f)
+    # phase_f = np.angle(f)
+
+    # make generic matrices, where r represents frequency,
+    # theta for orientation
+    x, y = np.meshgrid(np.arange(-size_tmp/2, size_tmp/2),
+                       np.arange(-size_tmp/2, size_tmp/2))
+    r = np.sqrt(np.power(x, 2) + np.power(y, 2))
+    y[y == 0] = .01
+    theta = np.arctan(x / y)
+    theta[f.shape[1] // 2:, :] = theta[f.shape[1] // 2 :, :] - np.pi
+    theta += 3 * np.pi / 2
+    # theta = np.arctan2(y, x) + np.pi  #in radians
+
+    # build the filter, spatial freq filters
+    if if_low_pass:
+        if np.isinf(sf_sigma):
+            sf_band = np.zeros(r.shape)
+            sf_band[r <= max_sf] = 1
+        else:
+            sf_band = np.exp(-((r - center_sf) ** 2 / 2 / sf_sigma ** 2))
+    elif if_band_pass:
+        sf_band = np.zeros(r.shape)
+        sf_band[(r >= low_sf) & (r<= high_sf)] = 1
+
+
+    # orientation filters
+    if np.isinf(orient_sigma):
+        pass_band = sf_band
+    else:  # need to make sure it works properly for 0-180 deg
+        orient_band = np.exp(-((theta - (orientation + np.pi / 2)) ** 2 / 2 / orient_sigma ** 2)) + \
+                      np.exp(-((np.angle(np.exp(1j * (theta))) + np.pi - (orientation + np.pi / 2)) ** 2
+                               / 2 / orient_sigma ** 2))
+        pass_band = np.multiply(orient_band, sf_band)
+    # only pass the needed components, and reconstruct the image back
+    band = np.fft.fftshift(np.fft.ifft2(np.fft.fftshift(np.multiply(f, pass_band))))
+    band = np.real(band)
+
+    # normalize for contrast
+    band_min = np.min(band[:])
+    band_max = np.max(band[:])
+    band = (band - band_min) * contrast * 2 / (band_max - band_min) - 1
+
+    # crop if needed
+    if size_tmp > size:
+        band = band[0:size, 0:size]
+
+    return band
+