From 78ddd059dfed5c73307fa1cef52fabaf40127cb1 Mon Sep 17 00:00:00 2001 From: gdamms Date: Tue, 31 Jan 2023 18:06:15 +0100 Subject: [PATCH] generate img with cv2 not plt --- src/draw.py | 73 ++++++++++++++++++++++++++++------------------------- 1 file changed, 38 insertions(+), 35 deletions(-) diff --git a/src/draw.py b/src/draw.py index d41a4d8..36ec1b0 100644 --- a/src/draw.py +++ b/src/draw.py @@ -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)