A holiday card for 2019

holiday19.py

import io
import requests
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap as LSC
from scipy.interpolate import Rbf

class HolidayCard():
    """A holiday card class."""
    def __init__(self, n_samples=1000, sn=100, seed=42):
        np.random.seed(seed)
        self.n_samples = n_samples
        self.sn = sn
        return

    def sample(self, fname):
        """Sample the image according to hyperparameters."""
        if fname.startswith("http"):
            data = requests.get(fname).content
            fname = io.BytesIO(data)
        img = Image.open(fname).convert('L')
        arr = np.asarray(img) / 255
        arr = np.flipud(arr)
        arr += np.random.normal(loc=0.5,
                                size=arr.shape
                               ) / self.sn
        h, w = arr.shape
        y, x = np.array([[np.random.randint(0, h),
                          np.random.randint(0, w)]
                         for _ in range(self.n_samples)]).T
        self.y, self.x = y, x
        self.z = arr[y, x]
        return

    def make_grid(self, ds=10, pad=0.01):
        """Make a grid for interpolation."""
        self.ds = ds
        xn, xx = self.x.min(), self.x.max()
        yn, yx = self.y.min(), self.y.max()
        e = max(xx-xn, yx-yn) * pad
        self.extent = (xn - e, xx + e, yn - e, yx + e)
        return np.mgrid[xn-e : xx+e : ds,
                        yn-e : yx+e : ds]

    def interpolate(self, grid):
        """Interpolate with radial basis function."""
        x_grid, y_grid = grid
        rbfi = Rbf(self.x, self.y, self.z, smooth=0.02)
        self.z_grid = rbfi(x_grid, y_grid)
        return

    def plot(self, interval=0.05, out=None, cmap=None, snow=True, c='w'):
        """Make the card!"""
        mi, ma = self.z_grid.min(), self.z_grid.max()
        levels = np.arange(mi, ma+0.2, interval)
        if cmap is None:
            cmap = LSC.from_list("", ["limegreen", "red"])

        fig, ax = plt.subplots(figsize=(8, 8))

        ax.contourf(self.z_grid.T,
                    cmap=cmap,
                    levels=levels,
                    extent=self.extent)

        ax.contour(self.z_grid.T,
                   levels=levels,
                   colors=[c],
                   linewidths=[1.5],
                   alpha=0.75,
                   extent=self.extent)
    
        if snow:
            alpha = np.random.random(self.x.size).reshape(-1, 1)
            c = np.hstack([np.ones((self.x.size, 3)), alpha])
            ax.scatter(self.y, self.x, c=c, marker='o', s=10)

        ax.axis('off')

        plt.tight_layout()
        plt.savefig(out or "output.png", dpi=250)
        plt.show()
        return


if __name__ == "__main__":

    card = HolidayCard(n_samples=1000, seed=2019, sn=25)
    card.sample("https://ageo.co/XVNNv8")
    grid = card.make_grid()
    card.interpolate(grid)
    card.plot(interval=0.1, out="xmas_2019.pdf")