tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Added video cropping for equirectangular and made it independent of mask in normal videos (#1542)

Browse Source 
* started bottom crop functionality

* first draft of cropping

* cleaned up and debugged cropping equirectangular and re added mask stuff just without percent crop

* added nonequirectangular video crop

* added left right crop for equirectangular

* shortened equirect utils

* finished regular video cropping

* removed print, added underscores, changed type to List[float]
This commit is contained in:
Rohan Mathur 2023-03-02 18:24:59 -08:00 committed by GitHub
parent f48fa6a923
commit 982e527bee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 135 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

116
nerfstudio/process_data/equirect_utils.py
View File

@ -16,7 +16,7 @@
import os
from pathlib import Path
from typing import Tuple
from typing import List, Tuple
import cv2
import numpy as np
@ -33,10 +33,83 @@ from rich.progress import (
from nerfstudio.utils.rich_utils import ItersPerSecColumn
def _crop_bottom(bound_arr: list, fov: int, percent_crop: float) -> List[float]:
"""Returns a list of vertical bounds with the bottom cropped.
Args:
bound_arr (list): List of vertical bounds in ascending order.
fov (int): Field of view of the camera.
percent_crop (float): Percent of the image to crop from the bottom.
Returns:
list: A new list of bounds with the bottom cropped.
"""
degrees_chopped = 180 * percent_crop
new_bottom_start = 90 - degrees_chopped - fov / 2
for i, el in reversed(list(enumerate(bound_arr))):
if el > new_bottom_start + fov / 2:
bound_arr[i] = None
elif el > new_bottom_start:
diff = el - new_bottom_start
bound_arr[i] = new_bottom_start
for j in range(i - 1, -1, -1):
bound_arr[j] -= diff / (2 ** (i - j))
break
return bound_arr
def _crop_top(bound_arr: list, fov: int, percent_crop: float) -> List[float]:
"""Returns a list of vertical bounds with the top cropped.
Args:
bound_arr (list): List of vertical bounds in ascending order.
fov (int): Field of view of the camera.
percent_crop (float): Percent of the image to crop from the top.
Returns:
list: A new list of bounds with the top cropped.
"""
degrees_chopped = 180 * percent_crop
new_top_start = -90 + degrees_chopped + fov / 2
for i, el in enumerate(bound_arr):
if el < new_top_start - fov / 2:
bound_arr[i] = None
elif el < new_top_start:
diff = new_top_start - el
bound_arr[i] = new_top_start
for j in range(i + 1, len(bound_arr)):
bound_arr[j] += diff / (2 ** (j - i))
break
return bound_arr
def _crop_bound_arr_vertical(
bound_arr: list, fov: int, percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0)
) -> list:
"""Returns a list of vertical bounds adjusted for cropping.
Args:
bound_arr (list): Original list of vertical bounds in ascending order.
fov (int): Field of view of the camera.
percent_crop (Tuple[float, float, float, float]): Crop arr (top, bottom, left, right).
Returns:
list: _description_
"""
if percent_crop[1] > 0:
bound_arr = _crop_bottom(bound_arr, fov, percent_crop[1])
if percent_crop[0] > 0:
bound_arr = _crop_top(bound_arr, fov, percent_crop[0])
return bound_arr
def generate_planar_projections_from_equirectangular(
image_dir: Path,
planar_image_size: Tuple[int, int],
samples_per_im: int,
percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
) -> Path:
"""Generate planar projections from an equirectangular image.
@ -44,6 +117,7 @@ def generate_planar_projections_from_equirectangular(
image_dir: The directory containing the equirectangular image.
planar_image_size: The size of the planar projections [width, height].
samples_per_im: The number of samples to take per image.
percent_crop: The percent of the image to crop from the (top, bottom, left, and right).
returns:
The path to the planar projections directory.
"""
@ -52,22 +126,38 @@ def generate_planar_projections_from_equirectangular(
fov = 120
yaw_pitch_pairs = []
left_bound, right_bound = -180, 180
if percent_crop[3] > 0:
left_bound = -180 + 360 * percent_crop[3]
if percent_crop[2] > 0:
right_bound = 180 - 360 * percent_crop[2]
if samples_per_im == 8:
fov = 120
for i in np.arange(-180, 180, 90):
yaw_pitch_pairs.append((i, 0))
for i in np.arange(-180, 180, 180):
yaw_pitch_pairs.append((i, 45))
for i in np.arange(-180, 180, 180):
yaw_pitch_pairs.append((i, -45))
bound_arr = [-45, 0, 45]
bound_arr = _crop_bound_arr_vertical(bound_arr, fov, percent_crop)
if bound_arr[1] is not None:
for i in np.arange(left_bound, right_bound, 90):
yaw_pitch_pairs.append((i, bound_arr[1]))
if bound_arr[2] is not None:
for i in np.arange(left_bound, right_bound, 180):
yaw_pitch_pairs.append((i, bound_arr[2]))
if bound_arr[0] is not None:
for i in np.arange(left_bound, right_bound, 180):
yaw_pitch_pairs.append((i, bound_arr[0]))
elif samples_per_im == 14:
fov = 110
for i in np.arange(-180, 180, 60):
yaw_pitch_pairs.append((i, 0))
for i in np.arange(-180, 180, 90):
yaw_pitch_pairs.append((i, 45))
for i in np.arange(-180, 180, 90):
yaw_pitch_pairs.append((i, -45))
bound_arr = [-45, 0, 45]
bound_arr = _crop_bound_arr_vertical(bound_arr, fov, percent_crop)
if bound_arr[1] is not None:
for i in np.arange(left_bound, right_bound, 60):
yaw_pitch_pairs.append((i, bound_arr[1]))
if bound_arr[2] is not None:
for i in np.arange(left_bound, right_bound, 90):
yaw_pitch_pairs.append((i, bound_arr[2]))
if bound_arr[0] is not None:
for i in np.arange(left_bound, right_bound, 90):
yaw_pitch_pairs.append((i, bound_arr[0]))
equi2pers = Equi2Pers(height=planar_image_size[1], width=planar_image_size[0], fov_x=fov, mode="bilinear")
frame_dir = image_dir

20
nerfstudio/process_data/process_data_utils.py
View File

@ -100,7 +100,11 @@ def get_num_frames_in_video(video: Path) -> int:
def convert_video_to_images(
video_path: Path, image_dir: Path, num_frames_target: int, verbose: bool = False
video_path: Path,
image_dir: Path,
num_frames_target: int,
percent_crop: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
verbose: bool = False,
) -> Tuple[List[str], int]:
"""Converts a video into a sequence of images.
@ -108,6 +112,7 @@ def convert_video_to_images(
video_path: Path to the video.
output_dir: Path to the output directory.
num_frames_target: Number of frames to extract.
percent_crop: Percent of the image to crop. (top, bottom, left, right)
verbose: If True, logs the output of the command.
Returns:
A tuple containing summary of the conversion and the number of extracted frames.
@ -135,16 +140,23 @@ def convert_video_to_images(
out_filename = image_dir / "frame_%05d.png"
ffmpeg_cmd = f'ffmpeg -i "{video_path}"'
spacing = num_frames // num_frames_target
crop_cmd = ""
if percent_crop != (0.0, 0.0, 0.0, 0.0):
height = 1 - percent_crop[0] - percent_crop[1]
width = 1 - percent_crop[2] - percent_crop[3]
start_x = percent_crop[2]
start_y = percent_crop[0]
crop_cmd = f',"crop=w=iw*{width}:h=ih*{height}:x=iw*{start_x}:y=ih*{start_y}"'
spacing = num_frames // num_frames_target
if spacing > 1:
ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB -r 1"
ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB{crop_cmd} -r 1"
else:
CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.")
ffmpeg_cmd += " -pix_fmt bgr8"
ffmpeg_cmd += f" {out_filename}"
run_command(ffmpeg_cmd, verbose=verbose)
num_final_frames = len(list(image_dir.glob("*.png")))

21
scripts/process_data.py
View File

@ -319,11 +319,19 @@ class ProcessVideo:
temp_image_dir = self.output_dir / "temp_images"
temp_image_dir.mkdir(parents=True, exist_ok=True)
summary_log, num_extracted_frames = process_data_utils.convert_video_to_images(
self.data, image_dir=temp_image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose
self.data,
image_dir=temp_image_dir,
num_frames_target=self.num_frames_target,
percent_crop=(0.0, 0.0, 0.0, 0.0),
verbose=self.verbose,
)
else:
summary_log, num_extracted_frames = process_data_utils.convert_video_to_images(
self.data, image_dir=image_dir, num_frames_target=self.num_frames_target, verbose=self.verbose
self.data,
image_dir=image_dir,
num_frames_target=self.num_frames_target,
percent_crop=self.percent_crop,
verbose=self.verbose,
)
# Generate planar projections if equirectangular
@ -332,7 +340,10 @@ class ProcessVideo:
self.output_dir / "temp_images", self.images_per_equirect
)
image_dir = equirect_utils.generate_planar_projections_from_equirectangular(
self.output_dir / "temp_images", perspective_image_size, self.images_per_equirect
self.output_dir / "temp_images",
perspective_image_size,
self.images_per_equirect,
percent_crop=self.percent_crop,
)
# copy the perspective images to the image directory
@ -346,11 +357,11 @@ class ProcessVideo:
# remove the temp_images folder
shutil.rmtree(self.output_dir / "temp_images", ignore_errors=True)
# Create mask
# # Create mask
mask_path = process_data_utils.save_mask(
image_dir=image_dir,
num_downscales=self.num_downscales,
percent_crop=self.percent_crop,
percent_crop=(0.0, 0.0, 0.0, 0.0),
percent_radius=self.percent_radius_crop,
)
if mask_path is not None: