nodes/wTechArtist-ComfyUI_VVL_VideoCamera_Advanced
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add function documents; re-format

Browse Source
This commit is contained in:
jianyuan
2025-03-17 23:31:23 +00:00
parent f1bc88b9f0
commit 56e693172d
6 changed files with 190 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
demo_gradio.py
+90 -11
View File
@@ -217,7 +217,7 @@ def gradio_demo(
# Handle None frame_filter
if frame_filter is None:
frame_filter = "All"
# Build a GLB file name
glbfile = os.path.join(
target_dir,
@@ -379,9 +379,9 @@ with gr.Blocks(
# Instead of gr.State, we use a hidden Textbox:
is_example = gr.Textbox(label="is_example", visible=False, value="None")
num_images = gr.Textbox(label="num_images", visible=False, value="None")
gr.HTML(
"""
"""
<h1>🏛️ VGGT: Visual Geometry Grounded Transformer</h1>
<p>
<a href="https://github.com/facebookresearch/vggt">🐙 GitHub Repository</a> |
@@ -542,7 +542,16 @@ with gr.Blocks(
fn=update_log, inputs=[], outputs=[log_output]
).then(
fn=gradio_demo,
inputs=[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode],
inputs=[
target_dir_output,
conf_thres,
frame_filter,
mask_black_bg,
mask_white_bg,
show_cam,
mask_sky,
prediction_mode,
],
outputs=[reconstruction_output, log_output, frame_filter],
).then(
fn=lambda: "False", inputs=[], outputs=[is_example] # set is_example to "False"
@@ -553,37 +562,107 @@ with gr.Blocks(
# -------------------------------------------------------------------------
conf_thres.change(
update_visualization,
[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example],
[
target_dir_output,
conf_thres,
frame_filter,
mask_black_bg,
mask_white_bg,
show_cam,
mask_sky,
prediction_mode,
is_example,
],
[reconstruction_output, log_output],
)
frame_filter.change(
update_visualization,
[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example],
[
target_dir_output,
conf_thres,
frame_filter,
mask_black_bg,
mask_white_bg,
show_cam,
mask_sky,
prediction_mode,
is_example,
],
[reconstruction_output, log_output],
)
mask_black_bg.change(
update_visualization,
[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example],
[
target_dir_output,
conf_thres,
frame_filter,
mask_black_bg,
mask_white_bg,
show_cam,
mask_sky,
prediction_mode,
is_example,
],
[reconstruction_output, log_output],
)
mask_white_bg.change(
update_visualization,
[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example],
[
target_dir_output,
conf_thres,
frame_filter,
mask_black_bg,
mask_white_bg,
show_cam,
mask_sky,
prediction_mode,
is_example,
],
[reconstruction_output, log_output],
)
show_cam.change(
update_visualization,
[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example],
[
target_dir_output,
conf_thres,
frame_filter,
mask_black_bg,
mask_white_bg,
show_cam,
mask_sky,
prediction_mode,
is_example,
],
[reconstruction_output, log_output],
)
mask_sky.change(
update_visualization,
[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example],
[
target_dir_output,
conf_thres,
frame_filter,
mask_black_bg,
mask_white_bg,
show_cam,
mask_sky,
prediction_mode,
is_example,
],
[reconstruction_output, log_output],
)
prediction_mode.change(
update_visualization,
[target_dir_output, conf_thres, frame_filter, mask_black_bg, mask_white_bg, show_cam, mask_sky, prediction_mode, is_example],
[
target_dir_output,
conf_thres,
frame_filter,
mask_black_bg,
mask_white_bg,
show_cam,
mask_sky,
prediction_mode,
is_example,
],
[reconstruction_output, log_output],
)
demo_viser.py
+18 -24
View File
@@ -30,6 +30,7 @@ from vggt.utils.load_fn import load_and_preprocess_images
from vggt.utils.geometry import closed_form_inverse_se3, unproject_depth_map_to_point_map
from vggt.utils.pose_enc import pose_encoding_to_extri_intri
def viser_wrapper(
pred_dict: dict,
port: int = 8080,
@@ -214,9 +215,9 @@ def viser_wrapper(
# Here we compute the threshold value based on the current percentage
current_percentage = gui_points_conf.value
threshold_val = np.percentile(conf_flat, current_percentage)
print(f"Threshold absolute value: {threshold_val}, percentage: {current_percentage}%")
conf_mask = (conf_flat >= threshold_val) & (conf_flat > 1e-5)
if gui_frame_selector.value == "All":
@@ -280,51 +281,44 @@ def apply_sky_segmentation(conf: np.ndarray, image_folder: str) -> np.ndarray:
np.ndarray: Updated confidence scores with sky regions masked out
"""
S, H, W = conf.shape
sky_masks_dir = image_folder.rstrip('/') + "_sky_masks"
sky_masks_dir = image_folder.rstrip("/") + "_sky_masks"
os.makedirs(sky_masks_dir, exist_ok=True)
# Download skyseg.onnx if it doesn't exist
if not os.path.exists("skyseg.onnx"):
print("Downloading skyseg.onnx...")
download_file_from_url(
"https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", "skyseg.onnx"
)
skyseg_session = onnxruntime.InferenceSession("skyseg.onnx")
download_file_from_url("https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", "skyseg.onnx")
skyseg_session = onnxruntime.InferenceSession("skyseg.onnx")
image_files = sorted(glob.glob(os.path.join(image_folder, "*")))
sky_mask_list = []
print("Generating sky masks...")
for i, image_path in enumerate(tqdm(image_files[:S])): # Limit to the number of images in the batch
image_name = os.path.basename(image_path)
mask_filepath = os.path.join(sky_masks_dir, image_name)
if os.path.exists(mask_filepath):
sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE)
else:
sky_mask = segment_sky(image_path, skyseg_session, mask_filepath)
# Resize mask to match H×W if needed
if sky_mask.shape[0] != H or sky_mask.shape[1] != W:
sky_mask = cv2.resize(sky_mask, (W, H))
sky_mask_list.append(sky_mask)
# Convert list to numpy array with shape S×H×W
sky_mask_array = np.array(sky_mask_list)
# Apply sky mask to confidence scores
sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32)
conf = conf * sky_mask_binary
print("Sky segmentation applied successfully")
return conf
parser = argparse.ArgumentParser(description="VGGT demo with viser for 3D visualization")
parser.add_argument(
"--image_folder", type=str, default="examples/kitchen/images/", help="Path to folder containing images"
@@ -341,14 +335,14 @@ parser.add_argument("--mask_sky", action="store_true", help="Apply sky segmentat
def main():
"""
Main function for the VGGT demo with viser for 3D visualization.
This function:
1. Loads the VGGT model
2. Processes input images from the specified folder
3. Runs inference to generate 3D points and camera poses
4. Optionally applies sky segmentation to filter out sky points
5. Visualizes the results using viser
Command-line arguments:
--image_folder: Path to folder containing input images
--use_point_map: Use point map instead of depth-based points
@@ -363,7 +357,7 @@ def main():
print("Initializing and loading VGGT model...")
# model = VGGT.from_pretrained("facebook/VGGT-1B")
model = VGGT()
_URL = "https://huggingface.co/facebook/VGGT-1B/resolve/main/model.pt"
model.load_state_dict(torch.hub.load_state_dict_from_url(_URL))
@@ -417,4 +411,4 @@ def main():
if __name__ == "__main__":
main()
main()
vggt/models/vggt.py
+15
View File
@@ -38,7 +38,22 @@ class VGGT(nn.Module, PyTorchModelHubMixin):
query_points (torch.Tensor, optional): Query points for tracking, in pixel coordinates.
Shape: [N, 2] or [B, N, 2], where N is the number of query points.
Default: None
Returns:
dict: A dictionary containing the following predictions:
- pose_enc (torch.Tensor): Camera pose encoding with shape [B, S, 9] (from the last iteration)
- depth (torch.Tensor): Predicted depth maps with shape [B, S, H, W, 1]
- depth_conf (torch.Tensor): Confidence scores for depth predictions with shape [B, S, H, W]
- world_points (torch.Tensor): 3D world coordinates for each pixel with shape [B, S, H, W, 3]
- world_points_conf (torch.Tensor): Confidence scores for world points with shape [B, S, H, W]
- images (torch.Tensor): Original input images, preserved for visualization
If query_points is provided, also includes:
- track (torch.Tensor): Point tracks with shape [B, S, N, 2] (from the last iteration), in pixel coordinates
- vis (torch.Tensor): Visibility scores for tracked points with shape [B, S, N]
- conf (torch.Tensor): Confidence scores for tracked points with shape [B, S, N]
"""
# If without batch dimension, add it
if len(images.shape) == 4:
images = images.unsqueeze(0)
vggt/utils/load_fn.py
+2 -2
View File
@@ -44,9 +44,9 @@ def load_and_preprocess_images(image_path_list):
img = Image.open(image_path)
# If there's an alpha channel, blend onto white background:
if img.mode == 'RGBA':
if img.mode == "RGBA":
# Create white background
background = Image.new('RGBA', img.size, (255, 255, 255, 255))
background = Image.new("RGBA", img.size, (255, 255, 255, 255))
# Alpha composite onto the white background
img = Image.alpha_composite(background, img)
vggt/utils/pose_enc.py
+62 -5
View File
@@ -13,9 +13,35 @@ def extri_intri_to_pose_encoding(
intrinsics,
image_size_hw=None, # e.g., (256, 512)
pose_encoding_type="absT_quaR_FoV",
min_focal_length=0.1,
max_focal_length=10,
):
"""Convert camera extrinsics and intrinsics to a compact pose encoding.
This function transforms camera parameters into a unified pose encoding format,
which can be used for various downstream tasks like pose prediction or representation.
Args:
extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4,
where B is batch size and S is sequence length.
In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world transformation.
The format is [R|t] where R is a 3x3 rotation matrix and t is a 3x1 translation vector.
intrinsics (torch.Tensor): Camera intrinsic parameters with shape BxSx3x3.
Defined in pixels, with format:
[[fx, 0, cx],
[0, fy, cy],
[0, 0, 1]]
where fx, fy are focal lengths and (cx, cy) is the principal point
image_size_hw (tuple): Tuple of (height, width) of the image in pixels.
Required for computing field of view values. For example: (256, 512).
pose_encoding_type (str): Type of pose encoding to use. Currently only
supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view).
Returns:
torch.Tensor: Encoded camera pose parameters with shape BxSx9.
For "absT_quaR_FoV" type, the 9 dimensions are:
- [:3] = absolute translation vector T (3D)
- [3:7] = rotation as quaternion quat (4D)
- [7:] = field of view (2D)
"""
# extrinsics: BxSx3x4
# intrinsics: BxSx3x3
@@ -25,7 +51,6 @@ def extri_intri_to_pose_encoding(
T = extrinsics[:, :, :3, 3] # BxSx3
quat = mat_to_quat(R)
# R_reverse = quat_to_mat(quat)
# Note the order of h and w here
H, W = image_size_hw
fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1])
@@ -40,11 +65,43 @@ def extri_intri_to_pose_encoding(
def pose_encoding_to_extri_intri(
pose_encoding,
image_size_hw=None, # e.g., (256, 512)
min_focal_length=0.1,
max_focal_length=10,
pose_encoding_type="absT_quaR_FoV",
build_intrinsics=True,
):
"""Convert a pose encoding back to camera extrinsics and intrinsics.
This function performs the inverse operation of extri_intri_to_pose_encoding,
reconstructing the full camera parameters from the compact encoding.
Args:
pose_encoding (torch.Tensor): Encoded camera pose parameters with shape BxSx9,
where B is batch size and S is sequence length.
For "absT_quaR_FoV" type, the 9 dimensions are:
- [:3] = absolute translation vector T (3D)
- [3:7] = rotation as quaternion quat (4D)
- [7:] = field of view (2D)
image_size_hw (tuple): Tuple of (height, width) of the image in pixels.
Required for reconstructing intrinsics from field of view values.
For example: (256, 512).
pose_encoding_type (str): Type of pose encoding used. Currently only
supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view).
build_intrinsics (bool): Whether to reconstruct the intrinsics matrix.
If False, only extrinsics are returned and intrinsics will be None.
Returns:
tuple: (extrinsics, intrinsics)
- extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4.
In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world
transformation. The format is [R|t] where R is a 3x3 rotation matrix and t is
a 3x1 translation vector.
- intrinsics (torch.Tensor or None): Camera intrinsic parameters with shape BxSx3x3,
or None if build_intrinsics is False. Defined in pixels, with format:
[[fx, 0, cx],
[0, fy, cy],
[0, 0, 1]]
where fx, fy are focal lengths and (cx, cy) is the principal point,
assumed to be at the center of the image (W/2, H/2).
"""
intrinsics = None
visual_util.py
+3 -3
View File
@@ -131,7 +131,7 @@ def predictions_to_glb(
sky_mask_array = np.array(sky_mask_list)
# Apply sky mask to confidence scores
sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32)
sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32)
pred_world_points_conf = pred_world_points_conf * sky_mask_binary
if selected_frame_idx is not None:
@@ -160,7 +160,7 @@ def predictions_to_glb(
if mask_black_bg:
black_bg_mask = colors_rgb.sum(axis=1) >= 16
conf_mask = conf_mask & black_bg_mask
if mask_white_bg:
# Filter out white background pixels (RGB values close to white)
# Consider pixels white if all RGB values are above 240
@@ -392,7 +392,7 @@ def segment_sky(image_path, onnx_session, mask_filename=None):
# The model outputs low values for sky, high values for non-sky
output_mask = np.zeros_like(result_map_original)
output_mask[result_map_original < 32] = 255 # Use threshold of 32
os.makedirs(os.path.dirname(mask_filename), exist_ok=True)
cv2.imwrite(mask_filename, output_mask)
return output_mask