generate img with cv2 not plt

This commit is contained in:
gdamms 2023-01-31 18:06:15 +01:00
parent aac57406a8
commit 78ddd059df

View file

@ -3,14 +3,20 @@ import numpy as np
import cv2
VOXEL_SIZE = 0.1
X_MIN, X_MAX = -5, 5
Y_MIN, Y_MAX = -5, 5
x_vals = np.arange(X_MIN, X_MAX, VOXEL_SIZE)
y_vals = np.arange(Y_MIN, Y_MAX, VOXEL_SIZE)
image_length = 1500
def f(x, y):
return np.exp(-((x**2) + y**2) / 3)
x_vals = np.linspace(-5, 5, 100)
y_vals = np.linspace(-5, 5, 100)
X, Y = np.meshgrid(x_vals, y_vals)
Z = f(X - 2, Y) + f(X + 2, Y)
@ -24,32 +30,36 @@ for i, x in enumerate(x_vals):
color = "#003"
else:
color = "#" + 3 * color
plt.fill([x, x + 0.1, x + 0.1, x], [y, y, y + 0.1, y + 0.1], color=color)
plt.fill([x, x + 0.1, x + 0.1, x],
[y, y, y + 0.1, y + 0.1], color=color)
nb_cams = 32
cam_poses = np.array(
[[6 * np.cos(theta), 6 * np.sin(theta)] for theta in np.linspace(0, 2 * np.pi, nb_cams, endpoint=False)]
[[6 * np.cos(theta), 6 * np.sin(theta)]
for theta in np.linspace(0, 2 * np.pi, nb_cams, endpoint=False)]
)
cam_rots = np.linspace(np.pi, 3 * np.pi, nb_cams, endpoint=False)
cam2world_projs = np.array(
[
[[np.cos(theta), -np.sin(theta), cam_pose[0]], [np.sin(theta), np.cos(theta), cam_pose[1]], [0, 0, 1]]
[[np.cos(theta), -np.sin(theta), cam_pose[0]],
[np.sin(theta), np.cos(theta), cam_pose[1]],
[0, 0, 1]]
for theta, cam_pose in zip(cam_rots, cam_poses)
]
)
for i in range(nb_cams):
plt.plot(cam_poses[i][0], cam_poses[i][1], "ro")
plt.text(cam_poses[i][0], cam_poses[i][1], str(i))
x = np.array([[0, 0, 1], [0.5, -0.2, 1], [0.5, 0.2, 1], [0, 0, 1]]).T
x = cam2world_projs[i] @ x
plt.plot(cam_poses[i][0], cam_poses[i][1], "ro")
plt.text(cam_poses[i][0], cam_poses[i][1], str(i))
plt.plot(x[0, :], x[1, :], "r-")
plt.xlim(-7, 7)
plt.ylim(-7, 7)
plt.axis("equal")
plt.savefig("data/peanut/peanut.png", dpi=300, bbox_inches="tight", pad_inches=0, transparent=True)
plt.savefig("data/peanut/peanut.png", dpi=300,
bbox_inches="tight", pad_inches=0, transparent=True)
plt.close()
# draw 1d image of the scene for each camera
@ -57,16 +67,15 @@ for i in range(nb_cams):
# sort pixels by distance to camera
cam_pose = cam_poses[i]
pixels_dist = np.linalg.norm(np.array([X.flatten(), Y.flatten()]).T - cam_pose, axis=1)
pixels_dist = np.linalg.norm(
np.array([X.flatten(), Y.flatten()]).T - cam_pose, axis=1)
pixels_sort = np.argsort(pixels_dist)[::-1]
x0 = -1
x1 = 1
plt.figure(f"img{i}")
plt.fill([x0, x1, x1, x0], [0, 0, 0.2, 0.2], color="#000")
plt.figure(f"mask{i}")
plt.fill([x0, x1, x1, x0], [0, 0, 0.2, 0.2], color="#000")
img = np.zeros((100, image_length, 3), dtype=np.uint8)
mask = np.zeros((100, image_length, 3), dtype=np.uint8)
for j in pixels_sort:
x, y = X.flatten()[j], Y.flatten()[j]
@ -75,28 +84,22 @@ for i in range(nb_cams):
continue
color = "#" + 3 * color
px = np.array([[x, y, 1], [x + 0.1, y, 1], [x + 0.1, y + 0.1, 1], [x, y + 0.1, 1]]).T
px = np.linalg.inv(cam2world_projs[i]) @ px
RT = np.linalg.inv(cam2world_projs[i])[:-1, :]
px = np.array([[x, y, 1],
[x + 0.1, y, 1],
[x + 0.1, y + 0.1, 1],
[x, y + 0.1, 1]])
px = RT @ px.T
px /= px[0, :]
px += np.array([[0], [1.0]])
px *= 0.5 * np.array([[1], [image_length]])
x0 = px[1, :].min()
x1 = px[1, :].max()
plt.figure(f"img{i}")
plt.fill([x0, x1, x1, x0], [0, 0, 0.2, 0.2], color=color)
plt.figure(f"mask{i}")
plt.fill([x0, x1, x1, x0], [0, 0, 0.2, 0.2], color="#fff")
img[:, int(x0):int(x1), :] = np.array(
3 * [int(Z.flatten()[j] * 255)], dtype=np.uint8)
mask[:, int(x0):int(x1), :] = np.array([255, 255, 255], dtype=np.uint8)
plt.figure(f"img{i}")
plt.xlim(-1, 1)
plt.ylim(0, 0.2)
plt.axis("off")
plt.savefig(f"data/peanut/images/Image{i:04}.png", dpi=300, bbox_inches="tight", pad_inches=0, transparent=True)
plt.close()
plt.figure(f"mask{i}")
plt.xlim(-1, 1)
plt.ylim(0, 0.2)
plt.axis("off")
plt.savefig(f"data/peanut/masks/Image{i:04}.png", dpi=300, bbox_inches="tight", pad_inches=0, transparent=True)
plt.close()
cv2.imwrite(f"data/peanut/images/Image{i:04}.png", img)
cv2.imwrite(f"data/peanut/masks/Image{i:04}.png", mask)