ReshotAI/gaussian-painters

Use more than three photo's for stenography experiment

arlorostirolla opened this issue · 1 comments

Hello. I would like to use 5 photo's with the stenography experiment. Is this possible? If so, would you be able to give me a general overview of how i could modify the code to do this?

Figured it out. If anyone is interested in doing this here is the code:

def create_dataset(args):
data_dir = Path(args.output_dir)
data_dir.mkdir(exist_ok=True)

images_dir = data_dir / "images"
images_dir.mkdir(exist_ok=True)

sparse_dir = data_dir / "sparse" / "0"
sparse_dir.mkdir(exist_ok=True, parents=True)

img_path = Path(args.img_path)
img_path2 = Path(args.img_path2)
img_path3 = Path(args.img_path3)
img_path4 = Path(args.img_path4)
img_path5 = Path(args.img_path5)

shutil.copy2(img_path, images_dir)
shutil.copy2(img_path2, images_dir)
shutil.copy2(img_path3, images_dir)
shutil.copy2(img_path4, images_dir)
shutil.copy2(img_path5, images_dir)

img = Image.open(args.img_path)
width, height = img.size
focal = math.sqrt(width**2 + height**2)

img2 = Image.open(args.img_path2)
width2, height2 = img2.size
focal2 = math.sqrt(width2**2 + height2**2)

img3 = Image.open(args.img_path3)
width3, height3 = img3.size
focal3 = math.sqrt(width3**2 + height3**2)

img4 = Image.open(args.img_path4)
width4, height4 = img4.size
focal4 = math.sqrt(width4**2 + height4**2)

img5 = Image.open(args.img_path5)
width5, height5 = img5.size
focal5 = math.sqrt(width5**2 + height5**2)

# create black image
black_img = Image.new("RGB", (width, height), (0, 0, 0))
black_img.save(images_dir / "black.jpg")


with open(sparse_dir / "cameras.txt", "w") as f:
    f.write(f"1 PINHOLE {width} {height} {focal} {focal} {width/2} {height/2}\n")
    f.write(f"2 PINHOLE {width2} {height2} {focal2} {focal2} {width2/2} {height2/2}\n")
    f.write(f"3 PINHOLE {width3} {height3} {focal3} {focal3} {width3/2} {height3/2}\n")
    f.write(f"4 PINHOLE {width4} {height4} {focal4} {focal4} {width4/2} {height4/2}\n")
    f.write(f"5 PINHOLE {width5} {height5} {focal5} {focal5} {width5/2} {height5/2}\n")

with open(sparse_dir / "images.txt", "w") as f:
    num_cameras = 5  
    full_circle = 2 * math.pi  # 360 degrees in radians
    focals = [focal, focal2, focal3, focal4, focal5]
    img_paths = [img_path, img_path2, img_path3, img_path4, img_path5]
    
    for i in range(1, num_cameras + 1):
        # Calculate the rotation angle for each camera
        theta = (i - 1) * full_circle / num_cameras

        # Quaternion calculation
        qx = 0
        qy = math.sin(theta / 2)
        qz = 0
        qw = math.cos(theta / 2)

        with open(sparse_dir / "images.txt", "a") as f:
            f.write(f"{i} {qx} {qy} {qz} {qw} 0 0 {focals[i-1]} {i} {img_paths[i-1].name}\n\n")

with open(sparse_dir / "points3D.txt", "w") as f:
    BLOCKS = 10
    idx = 0
    for i in range(0, BLOCKS):
        for j in range(0, BLOCKS):
            idx += 1
            x = i / BLOCKS * width
            y = j / BLOCKS * height
            rgb = img.getpixel((x, y))
            xyz = np.array([-width / 2, -height / 2, 0]) + np.array([x, y, 0])
            f.write(f"{idx} {' '.join(map(str, xyz.tolist()))} {' '.join(map(str, rgb))} 0\n")

subprocess.run(["colmap", "model_converter", "--input_path", sparse_dir, "--output_path", sparse_dir, "--output_type", "BIN"])