Spaces:
Runtime error
Runtime error
File size: 2,527 Bytes
c5d7102 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.tree import DecisionTreeRegressor
import PIL
def img2df(img_array):
dim0_arr = (
np.arange(img_array.shape[0])
.reshape((-1, 1))
.repeat(img_array.shape[1], axis=1)
.flatten()
)
dim1_arr = (
np.arange(img_array.shape[1])
.reshape((1, -1))
.repeat(img_array.shape[0], axis=0)
.flatten()
)
df = pd.DataFrame({"dim0": dim0_arr, "dim1": dim1_arr})
values = img_array.reshape((img_array.shape[0] * img_array.shape[1], -1))
for col in range(values.shape[1]):
df[f"value{col}"] = values[:, col]
return df
def normalize(img):
return (img - img.min()) / (img.max() - img.min())
def df2xy(df):
x = df[[c for c in df if c.startswith("dim")]]
y = df[[c for c in df if c.startswith("value")]]
if len(y.columns) == 1:
y = y.values.reshape(-1)
return x, y
def tree_window(
img, add_cartesian=True, add_rotation=False, add_polar=False, depths=(2, 6)
):
df = img2df(img)
x_raw, y = df2xy(df)
sets = []
if add_cartesian:
sets.append(x_raw)
if add_rotation > 0:
# rotate
theta = np.radians(add_rotation)
c, s = np.cos(theta), np.sin(theta)
R = np.array(((c, -s), (s, c)))
xr = x_raw @ R
sets.append(xr)
if add_polar:
# polar
xp = x_raw.copy()
xp["dim0"] = np.sqrt(((x_raw - x_raw.mean()) ** 2).sum(axis=1))
xp["dim1"] = np.arctan2(x_raw["dim1"], x_raw["dim0"])
sets.append(xp)
x = pd.concat(sets, axis=1)
fig, axes = plt.subplots(ncols=len(depths) + 1, figsize=(36, 36))
for ax, depth in zip(axes, depths):
model = DecisionTreeRegressor(max_depth=depth)
model.fit(x, y)
pred = model.predict(x).reshape(img.shape)
if len(y.shape) == 2:
ax.imshow(pred)
else:
ax.imshow(pred, cmap="gray")
ax.set_axis_off()
if len(y.shape) == 2:
axes[-1].imshow(img)
else:
axes[-1].imshow(img, cmap="gray")
axes[-1].set_axis_off()
return fig
def treeify(img, max_depth):
img_array = np.array(img)
df = img2df(img_array)
x, y = df2xy(df)
model = DecisionTreeRegressor(max_depth=max_depth)
model.fit(x, y)
pred = PIL.Image.fromarray(
model.predict(x).reshape(img_array.shape).round().astype("uint8")
)
score = model.score(x, y)
return pred, score
|