Permalink
Cannot retrieve contributors at this time
""" | |
Copyright (c) 2018 Shane Steinert-Threlkeld | |
***** | |
This program is free software: you can redistribute it and/or modify | |
it under the terms of the GNU General Public License as published by | |
the Free Software Foundation, either version 3 of the License, or | |
(at your option) any later version. | |
This program is distributed in the hope that it will be useful, | |
but WITHOUT ANY WARRANTY; without even the implied warranty of | |
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the | |
GNU General Public License for more details. | |
You should have received a copy of the GNU General Public License | |
along with this program. If not, see <http://www.gnu.org/licenses/>. | |
***** | |
This module contains methods that generate images of colored dots | |
for use in psychological studies, especially those testing numerical abilities. | |
In particular, it contains Python implementations of image generators for | |
all four conditions of the study presented in the paper | |
* Pietroski, P., Lidz, J., Hunter, T., & Halberda, J. (2009). | |
The Meaning of `Most’: Semantics, Numerosity and Psychology. | |
_Mind and Language_, 24(5), 554–585. | |
Usage as follows: | |
* each trial type has a corresponding method, e.g. `scattered_pairs` | |
* these methods take a `colors_dict` argument, a dictionary whose keys | |
are strings corresponding to matplotlib colors and whose values are | |
the number of dots of that color, e.g. | |
{'y': 9, 'b': 10} | |
* each method returns a list of `Dot`, where a `Dot` is an object with | |
four fields: x, y, radius, and color | |
* the method `make_image` takes such a list of `Dot`s and then | |
generates and saves an image to disk | |
* the method `make_batch` takes a list of trial types, a list of of | |
`colors_dict`s, and a number N, and produces N images of each | |
trial_type, color_dict pair | |
""" | |
from collections import namedtuple | |
import random | |
import math | |
from matplotlib.patches import Circle | |
import matplotlib.pyplot as plt | |
Dot = namedtuple('Dot', ['x', 'y', 'radius', 'color']) | |
def clip(val, min_val, max_val): | |
"""Clips `val` to be in the range [min_val, max_val]. """ | |
return max(min(val, max_val), min_val) | |
def polar_to_cartesian(theta, r): | |
"""Converts polar coordinates to Cartesian. """ | |
return r*math.cos(theta), r*math.sin(theta) | |
def no_overlap(dots, x, y, radius): | |
"""Checks whether a new dot will have any overlap with an existing | |
array of dots. | |
Args: | |
dots: an iterable of `Dot`s | |
x: x-value of center of new dot | |
y: y-value of center of new dot | |
r: radius of new dot | |
Returns: | |
True if the new dot has no overlap with any of the dots in `dots', | |
False otherwise | |
""" | |
return all([(x - dot.x)**2 + (y - dot.y)**2 >= (radius + dot.radius)**2 | |
for dot in dots]) | |
def get_random_radii(colors_dict, min_radius, max_radius, std=1): | |
"""Gets random radii of dots for a color_dict. Radii are sampled from | |
a Gaussian distribution with mean (max_r - min_r) / 2 and standard | |
deviation std, then clipped. | |
Args: | |
color_dict: dictionary of colors, with integer values | |
min_radius: smallest radius | |
max_radius: biggest radius | |
std: standard deviation | |
Returns: | |
a dictionary, with the same keys as colors_dict, and values a list of | |
colors_dict[color] floating point numbers | |
""" | |
mean = (max_radius - min_radius) / 2 | |
return {color: [clip(random.gauss(mean, std), min_radius, max_radius) | |
for _ in range(colors_dict[color])] | |
for color in colors_dict} | |
def get_area_controlled_radii(colors_dict, min_radius, max_radius, std=0.5, | |
total_area=None): | |
"""Gets area controlled radii: the sum of the areas of circles of each | |
color will be equal (either to total_area or to the total area taken by the | |
largest number in colors_dict dots of mean radius). | |
Args: | |
colors_dict: as above | |
min_radius: as above | |
max_radius: as above | |
std: as above | |
total_area: a float, the total area to distribute to each color. If | |
not specified, this will be set to N*(max_radius - min_radius)/2^2, | |
where N is the largest value in colors_dict | |
Returns: | |
a dictionary, as above | |
""" | |
mean = (max_radius - min_radius) / 2 | |
if not total_area: | |
total_area = math.pi*(mean**2)*max(colors_dict.values()) | |
radii = {color: [] for color in colors_dict} | |
for color in colors_dict: | |
num_remaining = colors_dict[color] | |
area_remaining = total_area | |
while num_remaining > 1: | |
mean = math.sqrt(area_remaining / (num_remaining*math.pi)) | |
# get radius that is not too big to use up all remaining area! | |
found_r = False | |
while not found_r: | |
r = clip(random.gauss(mean, std), min_radius, max_radius) | |
if math.pi*r**2 < area_remaining: | |
found_r = True | |
radii[color].append(r) | |
area_remaining -= math.pi*r**2 | |
num_remaining -= 1 | |
radii[color].append(math.sqrt(area_remaining / math.pi)) | |
return radii | |
def scattered_random(colors_dict, area_control=False, | |
total_area=None, | |
num_pixels=(256, 256), padding=16, | |
min_radius=1, max_radius=5, std=1): | |
"""Generates ScatteredRandom images: the dots are scattered | |
randomly through the image. """ | |
x_min, y_min = padding, padding | |
x_max, y_max = num_pixels[0] - padding, num_pixels[1] - padding | |
dots = [] | |
if area_control: | |
radii = get_area_controlled_radii(colors_dict, min_radius, max_radius, | |
std=std, total_area=total_area) | |
else: | |
radii = get_random_radii(colors_dict, min_radius, max_radius, std=std) | |
# print({color: sum([math.pi*r**2 for r in radii[color]]) for color in radii}) | |
for color in colors_dict: | |
for r in radii[color]: | |
new_dot_added = False | |
while not new_dot_added: | |
x = random.uniform(x_min, x_max) | |
y = random.uniform(y_min, y_max) | |
# avoid overlap with existing circles | |
if no_overlap(dots, x, y, r): | |
dots.append(Dot(x=x, y=y, radius=r, color=color)) | |
new_dot_added = True | |
return dots | |
def scattered_split(colors_dict, area_control=False, | |
num_pixels=(512, 256), padding=24, | |
min_radius=1, max_radius=5, std=0.5, | |
color_order=None): | |
"""Generates ScatteredSplit images: the dots are scattered randomly through | |
the image, but each color has its own region of the image, with different | |
colors laid out horizontally. """ | |
width_per = num_pixels[0] / len(colors_dict) | |
mean = (max_radius - min_radius) / 2 | |
total_area = math.pi*(mean**2)*max(colors_dict.values()) | |
color_dots = {color: scattered_random( | |
{color: colors_dict[color]}, area_control=area_control, | |
total_area=total_area, | |
num_pixels=(width_per, num_pixels[1]), padding=padding, | |
min_radius=min_radius, max_radius=max_radius, std=std) | |
for color in colors_dict} | |
dots = [] | |
if not color_order: | |
colors = list(colors_dict.keys()) | |
random.shuffle(colors) | |
else: | |
colors = color_order | |
for idx in range(len(colors)): | |
dots.extend([dot._replace(x=dot.x + idx*width_per) | |
for dot in color_dots[colors[idx]]]) | |
return dots | |
# TODO: test the tuple num_pixels in the three following methods | |
# TODO: add area_control to these three methods a la scattered_random | |
def scattered_pairs(colors_dict, num_pixels=(256, 256), padding=16, | |
min_radius=1, max_radius=5): | |
"""Generates ScatteredPairs images: the dots are paired together, one of | |
each type of color. The remaining dots in the dominant color are randomly | |
scattered. | |
NOTE: `colors_dict` must have only two keys. """ | |
x_min, y_min = padding, padding | |
x_max, y_max = num_pixels[0] - padding, num_pixels[1] - padding | |
assert len(colors_dict) == 2 | |
sort_dict = sorted(colors_dict.items(), key=lambda pair: pair[1]) | |
num_leftover = sort_dict[1][1] - sort_dict[0][1] | |
dots = [] | |
for _ in range(sort_dict[0][1]): | |
new_pair_added = False | |
while not new_pair_added: | |
# one circle | |
x1 = random.uniform(x_min, x_max) | |
y1 = random.uniform(y_min, y_max) | |
r1 = clip(random.gauss(3, 1), min_radius, max_radius) | |
# get second dot coordinates | |
r2 = clip(random.gauss(3, 1), min_radius, max_radius) | |
theta = random.uniform(0, 2*math.pi) | |
dist = r1 + r2 | |
tmp_x2, tmp_y2 = polar_to_cartesian(theta, dist) | |
x2, y2 = tmp_x2 + x1, tmp_y2 + y1 | |
if no_overlap(dots, x1, y1, r1) and no_overlap(dots, x2, y2, r2): | |
dots.append(Dot(x=x1, y=y1, radius=r1, color=sort_dict[0][0])) | |
dots.append(Dot(x=x2, y=y2, radius=r2, color=sort_dict[1][0])) | |
new_pair_added = True | |
for _ in range(num_leftover): | |
new_dot_added = False | |
while not new_dot_added: | |
x = random.uniform(x_min, x_max) | |
y = random.uniform(y_min, y_max) | |
r = clip(random.gauss(3, 1), min_radius, max_radius) | |
# avoid overlap with existing circles | |
if no_overlap(dots, x, y, r): | |
dots.append(Dot(x=x, y=y, radius=r, color=sort_dict[1][0])) | |
new_dot_added = True | |
return dots | |
def column_pairs_mixed(colors_dict, num_pixels=(256, 256), pad=2.5, | |
min_radius=1, max_radius=5): | |
"""Generates ColumnPairsMixed images: the dots are paired in two columns, | |
but which color is in which column depends on the row. | |
NOTE: `colors_dict` must have only two keys. """ | |
assert len(colors_dict) == 2 | |
sort_dict = sorted(colors_dict.items(), key=lambda pair: pair[1]) | |
center = num_pixels[0] / 2 | |
x_vals = (center - max_radius - pad/2, center + max_radius + pad/2) | |
y_step = max_radius + 2*pad | |
y0 = center + int(sort_dict[1][1] / 2)*y_step | |
num_lower = sort_dict[0][1] | |
dots = [] | |
for row_num in range(sort_dict[1][1]): | |
first_side = random.choice([0, 1]) | |
x1 = x_vals[first_side] | |
y1 = y0 - row_num*y_step | |
r1 = clip(random.gauss(3, 1), min_radius, max_radius) | |
dots.append(Dot(x=x1, y=y1, radius=r1, color=sort_dict[1][0])) | |
add_second = False | |
if num_lower < sort_dict[1][1] - row_num: | |
if random.random() < 0.5: | |
add_second = True | |
else: | |
add_second = True | |
if add_second: | |
x2 = x_vals[0 if first_side == 1 else 1] | |
y2 = y1 | |
r2 = clip(random.gauss(3, 1), min_radius, max_radius) | |
dots.append(Dot(x=x2, y=y2, radius=r2, color=sort_dict[0][0])) | |
num_lower -= 1 | |
return dots | |
def column_pairs_sorted(colors_dict, num_pixels=(256, 256), pad=2.5, | |
min_radius=1, max_radius=5): | |
"""Generates ColumnPairsSorted images: the dots are paired in two columns, | |
with one color on the left and one on the right. | |
NOTE: `colors_dict` must have only two keys. """ | |
assert len(colors_dict) == 2 | |
sort_dict = sorted(colors_dict.items(), key=lambda pair: pair[1]) | |
center = num_pixels[0] / 2 | |
x_vals = (center - max_radius - pad/2, center + max_radius + pad/2) | |
y_step = max_radius + 2*pad | |
y0 = center + int(sort_dict[1][1] / 2)*y_step | |
num_lower = sort_dict[0][1] | |
lower_color_side = random.choice([0, 1]) | |
higher_color_side = 1 if lower_color_side == 0 else 0 | |
dots = [] | |
for row_num in range(sort_dict[1][1]): | |
x1 = x_vals[higher_color_side] | |
y1 = y0 - row_num*y_step | |
r1 = clip(random.gauss(3, 1), min_radius, max_radius) | |
dots.append(Dot(x=x1, y=y1, radius=r1, color=sort_dict[1][0])) | |
if row_num < num_lower: | |
x2 = x_vals[lower_color_side] | |
y2 = y1 | |
r2 = clip(random.gauss(3, 1), min_radius, max_radius) | |
dots.append(Dot(x=x2, y=y2, radius=r2, color=sort_dict[0][0])) | |
return dots | |
def make_image(file_name, dots, num_pixels=(256, 256)): | |
"""Make and save an image from a list of `Dot`s. | |
Args: | |
file_name: where to save the image | |
dots: an iterable of `Dot`s to draw | |
num_pixels: the saved image will be a square with num_pixels sides | |
""" | |
# TODO: update doc-string for proper width-height understanding! | |
plt.rcParams['axes.facecolor'] = 'grey' | |
plt.rcParams['axes.linewidth'] = 0 | |
dpi = min(num_pixels) | |
fig, ax = plt.subplots(dpi=dpi) | |
fig.set_size_inches((max(num_pixels) / dpi, 1)) | |
ax.set_xlim((0, num_pixels[0])) | |
ax.set_ylim((0, num_pixels[1])) | |
for dot in dots: | |
ax.add_patch( | |
Circle((dot.x, dot.y), radius=dot.radius, fc=dot.color)) | |
ax.set_xticks([]) | |
ax.set_yticks([]) | |
plt.tight_layout(pad=0) | |
# TODO: mkdir for file_name if not exists... | |
fig.savefig(file_name, facecolor='grey', pad_inches=0) | |
def make_batch(trial_types, color_dicts, num_per_dict, out_dir='.', | |
num_pixels=(256, 256), min_radius=1, max_radius=5, std=1, | |
area_control=True): | |
for trial_type in trial_types: | |
image_method = globals()[trial_type] | |
for dict_idx in range(len(color_dicts)): | |
color_dict = color_dicts[dict_idx] | |
for idx in range(num_per_dict): | |
make_image( | |
'{}/{}_{}_{}_{}.png'.format( | |
out_dir, trial_type, | |
'_'.join([str(key) + str(color_dict[key]) | |
for key in color_dict]), | |
idx, dict_idx), | |
image_method(color_dict, | |
num_pixels=num_pixels, | |
min_radius=min_radius, | |
max_radius=max_radius, std=std, | |
area_control=area_control), | |
num_pixels=num_pixels) | |
def make_split_batch(color_dicts, num_per_dict, out_dir='.', | |
num_pixels=(800, 400), min_radius=1.5, max_radius=3.5, | |
std=0.5, area_control=True): | |
for dict_idx in range(len(color_dicts)): | |
color_dict = color_dicts[dict_idx] | |
order = list(color_dict.keys()) | |
for idx in range(num_per_dict): | |
# TODO: integrate this order parameter into make_batch instead of | |
# having this as a separate method? | |
# balance order by flipping every time | |
order = list(reversed(order)) | |
make_image( | |
'{}/{}_{}_{}_{}.png'.format( | |
out_dir, 'scattered_split', | |
'_'.join([str(key) + str(color_dict[key]) | |
for key in color_dict]), | |
idx, dict_idx), | |
scattered_split(color_dict, | |
num_pixels=num_pixels, | |
min_radius=min_radius, | |
max_radius=max_radius, std=std, | |
color_order=order, area_control=area_control), | |
num_pixels=num_pixels) | |
def dicts_from_ratios(ratios, dicts_per_ratio, | |
colors=['b', 'y'], dot_range=(5, 25), multipliers=None): | |
dicts = [] | |
assert multipliers is None or len(multipliers) == len(ratios) | |
for idx in range(len(ratios)): | |
ratio = ratios[idx] | |
if not multipliers: | |
# TODO: better method of generating ratios? | |
mults = [n for n in range(1, 10) | |
if n*min(ratio) >= dot_range[0] | |
and n*max(ratio) <= dot_range[1]] | |
else: | |
mults = multipliers[idx] | |
for idx in range(dicts_per_ratio): | |
# reverse the ratio every trial to mix | |
ratio = list(reversed(ratio)) | |
# mult = random.choice(multipliers) | |
# TODO: make sure this logic works | |
mult = mults[idx % len(mults)] | |
dicts.append({colors[n]: mult*ratio[n] for n in range(len(colors))}) | |
return dicts | |
if __name__ == '__main__': | |
ratios = [(n, n+1) for n in range(1, 10)] | |
imgs_per_ratio = 100 | |
trial_types = ['scattered_random', 'scattered_pairs', | |
'column_pairs_mixed', 'column_pairs_sorted'] | |
color_dicts = dicts_from_ratios(ratios, imgs_per_ratio, dot_range=(5, 12)) | |
# TODO: refactor this to get proper # imgs per bin | |
# make training set | |
make_batch(trial_types, color_dicts, 5, 'images/small/train', | |
num_pixels=128, min_radius=2, max_radius=4) | |
# make val set | |
make_batch(trial_types, color_dicts, 1, 'images/small/val', | |
num_pixels=128, min_radius=2, max_radius=4) | |
# make test set | |
make_batch(trial_types, color_dicts, 1, 'images/small/test', | |
num_pixels=128, min_radius=2, max_radius=4) |