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

Last commit beofre following the convention of visual_util.py

Browse Source
This commit is contained in:
Atakan Topaloglu
2025-03-29 17:34:34 +03:00
parent daf7f94ba4
commit 009d8515d1
1 changed files with 271 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vggt_to_colmap.py
+271 -98
View File
@@ -7,6 +7,8 @@ import struct
from scipy.spatial.transform import Rotation
import sys
from PIL import Image
import cv2
import requests
# Add VGGT to path
sys.path.append("vggt/")
@@ -22,13 +24,15 @@ def load_model(device=None):
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
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))
model.eval()
model = model.to(device)
return model, device
def process_images(image_dir, model, device, conf_threshold=50.0):
def process_images(image_dir, model, device):
"""Process images with VGGT and return predictions."""
image_names = glob.glob(os.path.join(image_dir, "*"))
image_names = sorted([f for f in image_names if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
@@ -74,6 +78,7 @@ def process_images(image_dir, model, device, conf_threshold=50.0):
# Store original images for color extraction
predictions["original_images"] = original_images
predictions["images"] = np.array(original_images) / 255.0 # Add normalized images as in original code
return predictions, image_names
@@ -99,108 +104,266 @@ def extrinsic_to_colmap_format(extrinsics):
return np.array(quaternions), np.array(translations)
def is_sky_point(point, rgb):
"""Check if a point is likely to be sky based on position and color."""
# Check if point is above horizon (y is up in most 3D coordinates)
# and has blue-ish color (more blue than red or green)
return (point[1] > 1.0 and # High up
rgb[2] > 1.2 * rgb[0] and # More blue than red
rgb[2] > 1.2 * rgb[1]) # More blue than green
def download_file_from_url(url, filename):
"""Downloads a file from a URL, handling redirects."""
try:
# Get the redirect URL
response = requests.get(url, allow_redirects=False)
response.raise_for_status() # Raise HTTPError for bad requests (4xx or 5xx)
def is_black_bg_point(rgb, threshold=30):
"""Check if a point is likely a black background pixel."""
# Check if all RGB values are very low
return np.mean(rgb) < threshold
if response.status_code == 302: # Expecting a redirect
redirect_url = response.headers["Location"]
response = requests.get(redirect_url, stream=True)
response.raise_for_status()
else:
response = requests.get(url, stream=True)
response.raise_for_status()
def is_white_bg_point(rgb, threshold=230):
"""Check if a point is likely a white background pixel."""
# Check if all RGB values are very high
return np.mean(rgb) > threshold and np.std(rgb) < 15
with open(filename, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
print(f"Downloaded {filename} successfully.")
return True
except requests.exceptions.RequestException as e:
print(f"Error downloading file: {e}")
return False
def segment_sky(image_path, onnx_session, mask_filename=None):
"""
Segments sky from an image using an ONNX model.
"""
assert mask_filename is not None
image = cv2.imread(image_path)
result_map = run_skyseg(onnx_session, [320, 320], image)
# resize the result_map to the original image size
result_map_original = cv2.resize(result_map, (image.shape[1], image.shape[0]))
# Fix: Invert the mask so that 255 = non-sky, 0 = sky
# 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
def run_skyseg(onnx_session, input_size, image):
"""
Runs sky segmentation inference using ONNX model.
"""
import copy
# Pre process:Resize, BGR->RGB, Transpose, PyTorch standardization, float32 cast
temp_image = copy.deepcopy(image)
resize_image = cv2.resize(temp_image, dsize=(input_size[0], input_size[1]))
x = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB)
x = np.array(x, dtype=np.float32)
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
x = (x / 255 - mean) / std
x = x.transpose(2, 0, 1)
x = x.reshape(-1, 3, input_size[0], input_size[1]).astype("float32")
# Inference
input_name = onnx_session.get_inputs()[0].name
output_name = onnx_session.get_outputs()[0].name
onnx_result = onnx_session.run([output_name], {input_name: x})
# Post process
onnx_result = np.array(onnx_result).squeeze()
min_value = np.min(onnx_result)
max_value = np.max(onnx_result)
onnx_result = (onnx_result - min_value) / (max_value - min_value)
onnx_result *= 255
onnx_result = onnx_result.astype("uint8")
return onnx_result
def filter_and_prepare_points(predictions, conf_threshold, mask_sky=False, mask_black_bg=False,
mask_white_bg=False, stride=4):
mask_white_bg=False, stride=4, target_dir=None, prediction_mode="Depthmap and Camera Branch"):
"""
Filter points based on confidence and prepare for COLMAP format.
Uses stride to sample fewer points for efficiency.
Extracts RGB color from original images.
Applies optional filtering for sky and backgrounds.
Implementation matches the conventions in the original VGGT code.
"""
# Scale confidence threshold from percentage to actual value range
max_conf = np.max(predictions["depth_conf"])
conf_threshold = conf_threshold / 100.0 * max_conf
# Use the correct point cloud based on prediction mode
if "Pointmap" in prediction_mode:
print("Using Pointmap Branch")
if "world_points" in predictions:
pred_world_points = predictions["world_points"]
pred_world_points_conf = predictions.get("world_points_conf", np.ones_like(pred_world_points[..., 0]))
else:
print("Warning: world_points not found in predictions, falling back to depth-based points")
pred_world_points = predictions["world_points_from_depth"]
pred_world_points_conf = predictions.get("depth_conf", np.ones_like(pred_world_points[..., 0]))
else:
print("Using Depthmap and Camera Branch")
pred_world_points = predictions["world_points_from_depth"]
pred_world_points_conf = predictions.get("depth_conf", np.ones_like(pred_world_points[..., 0]))
print(f"Confidence threshold: {conf_threshold:.4f} (scaled from {conf_threshold*100/max_conf:.2f}%)")
print(f"Filtering options - Sky: {mask_sky}, Black BG: {mask_black_bg}, White BG: {mask_white_bg}")
# Get images for color extraction
colors_rgb = predictions["images"] # Already in normalized format (0-1)
colors_rgb = (colors_rgb * 255).astype(np.uint8)
# Initialize containers
points3D = []
point_indices = {} # Maps point hash to its index
image_points2D = [[] for _ in range(len(predictions["depth"]))]
# Get world points and confidences
world_points = predictions["world_points_from_depth"] # (num_images, H, W, 3)
confidences = predictions["depth_conf"] # (num_images, H, W)
original_images = predictions["original_images"] # List of numpy arrays (H, W, 3)
# Process each image
for img_idx in range(len(world_points)):
h, w = world_points[img_idx].shape[:2]
orig_img = original_images[img_idx]
# Handle possible resizing between original image and depth map
orig_h, orig_w = orig_img.shape[:2]
scale_y, scale_x = orig_h / h, orig_w / w
# Sample points with stride
for y in range(0, h, stride):
for x in range(0, w, stride):
if confidences[img_idx, y, x] > conf_threshold:
# Get 3D point
point3D = world_points[img_idx, y, x]
# Handle sky segmentation if requested
if mask_sky and target_dir is not None:
print("Applying sky segmentation mask")
try:
import onnxruntime
skyseg_session = None
target_dir_images = os.path.join(target_dir, "images")
os.makedirs(target_dir_images, exist_ok=True)
# Save the RGB images temporarily if they don't exist
image_list = []
for i, img in enumerate(colors_rgb):
img_path = os.path.join(target_dir_images, f"image_{i:04d}.png")
image_list.append(os.path.basename(img_path))
if not os.path.exists(img_path):
cv2.imwrite(img_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
# Ensure the output directory for masks exists
sky_masks_dir = os.path.join(target_dir, "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_success = download_file_from_url(
"https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx",
"skyseg.onnx"
)
if not download_success:
print("Failed to download skyseg model, skipping sky filtering")
mask_sky = False
if mask_sky: # Continue only if download succeeded
S, H, W = pred_world_points_conf.shape
sky_mask_list = []
for i, image_name in enumerate(image_list):
image_filepath = os.path.join(target_dir_images, image_name)
mask_filepath = os.path.join(sky_masks_dir, image_name)
# Skip invalid points
if not np.all(np.isfinite(point3D)):
continue
# Get RGB color from original image
# Map coordinates from depth map to original image
orig_y = min(int(y * scale_y), orig_h - 1)
orig_x = min(int(x * scale_x), orig_w - 1)
rgb = orig_img[orig_y, orig_x]
# Apply filters
if mask_sky and is_sky_point(point3D, rgb):
continue
if mask_black_bg and is_black_bg_point(rgb):
continue
if mask_white_bg and is_white_bg_point(rgb):
continue
# Generate a hash for this point
point_hash = hash_point(point3D, scale=100)
if point_hash not in point_indices:
# Create new point
point_idx = len(points3D)
point_indices[point_hash] = point_idx
point_entry = {
"id": point_idx,
"xyz": point3D,
"rgb": rgb,
"error": 1.0,
"track": [(img_idx, len(image_points2D[img_idx]))]
}
points3D.append(point_entry)
# Check if mask already exists
if os.path.exists(mask_filepath):
# Load existing mask
sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE)
else:
# Update existing point's track
point_idx = point_indices[point_hash]
points3D[point_idx]["track"].append((img_idx, len(image_points2D[img_idx])))
# Generate new mask
if skyseg_session is None:
skyseg_session = onnxruntime.InferenceSession("skyseg.onnx")
sky_mask = segment_sky(image_filepath, skyseg_session, mask_filepath)
# Add 2D point to image
image_points2D[img_idx].append((x, y, point_indices[point_hash]))
# 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)
pred_world_points_conf = pred_world_points_conf * sky_mask_binary
except (ImportError, Exception) as e:
print(f"Error in sky segmentation: {e}")
print("Falling back to basic sky filtering")
mask_sky = False
print(f"Generated {len(points3D)} 3D points with {sum(len(pts) for pts in image_points2D)} observations")
# Reshape for processing
vertices_3d = pred_world_points.reshape(-1, 3)
conf = pred_world_points_conf.reshape(-1)
colors_rgb_flat = colors_rgb.reshape(-1, 3)
# Apply confidence threshold exactly as in original code
if conf_threshold == 0.0:
conf_thres_value = 0.0
else:
conf_thres_value = np.percentile(conf, conf_threshold)
print(f"Using confidence threshold: {conf_threshold}% (value: {conf_thres_value:.4f})")
conf_mask = (conf >= conf_thres_value) & (conf > 1e-5)
# Apply black background filter if requested
if mask_black_bg:
print("Filtering black background")
black_bg_mask = colors_rgb_flat.sum(axis=1) >= 16
conf_mask = conf_mask & black_bg_mask
# Apply white background filter if requested
if mask_white_bg:
print("Filtering white background")
white_bg_mask = ~((colors_rgb_flat[:, 0] > 240) & (colors_rgb_flat[:, 1] > 240) & (colors_rgb_flat[:, 2] > 240))
conf_mask = conf_mask & white_bg_mask
# Filter points
filtered_vertices = vertices_3d[conf_mask]
filtered_colors = colors_rgb_flat[conf_mask]
# Check if we have any points left
if len(filtered_vertices) == 0:
print("Warning: No points remaining after filtering. Using default point.")
filtered_vertices = np.array([[0, 0, 0]])
filtered_colors = np.array([[200, 200, 200]])
print(f"Filtered to {len(filtered_vertices)} points")
# Now prepare the filtered points for COLMAP format
points3D = []
point_indices = {}
image_points2D = [[] for _ in range(len(pred_world_points))]
# Get the original 3D indices from the flat indices
S, H, W = pred_world_points.shape[0:3]
flat_indices = np.where(conf_mask)[0]
for flat_idx in flat_indices[::stride]: # Use stride to reduce point count
# Convert flat index back to 3D coordinates
img_idx = flat_idx // (H * W)
remainder = flat_idx % (H * W)
y = remainder // W
x = remainder % W
if img_idx >= S: # Handle edge cases
continue
point3D = vertices_3d[flat_idx]
rgb = filtered_colors[flat_idx - flat_indices[0]] if flat_idx >= flat_indices[0] else filtered_colors[0]
# Skip invalid points
if not np.all(np.isfinite(point3D)):
continue
# Generate a hash for this point to identify duplicates
point_hash = hash_point(point3D, scale=100)
if point_hash not in point_indices:
# Create new point
point_idx = len(points3D)
point_indices[point_hash] = point_idx
point_entry = {
"id": point_idx,
"xyz": point3D,
"rgb": rgb,
"error": 1.0,
"track": [(img_idx, len(image_points2D[img_idx]))]
}
points3D.append(point_entry)
else:
# Update existing point's track
point_idx = point_indices[point_hash]
points3D[point_idx]["track"].append((img_idx, len(image_points2D[img_idx])))
# Add 2D point to image
image_points2D[img_idx].append((x, y, point_indices[point_hash]))
print(f"Prepared {len(points3D)} 3D points with {sum(len(pts) for pts in image_points2D)} observations for COLMAP")
return points3D, image_points2D
def hash_point(point, scale=100):
@@ -378,7 +541,7 @@ def main():
parser.add_argument("--output_dir", type=str, default="colmap_output",
help="Directory to save COLMAP files")
parser.add_argument("--conf_threshold", type=float, default=50.0,
help="Confidence threshold (0-100) for including points")
help="Confidence threshold (0-100%) for including points")
parser.add_argument("--mask_sky", action="store_true",
help="Filter out points likely to be sky")
parser.add_argument("--mask_black_bg", action="store_true",
@@ -389,31 +552,41 @@ def main():
help="Output binary COLMAP files instead of text")
parser.add_argument("--stride", type=int, default=1,
help="Stride for point sampling (higher = fewer points)")
parser.add_argument("--prediction_mode", type=str, default="Depthmap and Camera Branch",
choices=["Depthmap and Camera Branch", "Pointmap Branch"],
help="Which prediction branch to use")
args = parser.parse_args()
# Create output directory
os.makedirs(args.output_dir, exist_ok=True)
# Create temporary directory for processing
temp_dir = os.path.join(args.output_dir, "temp")
os.makedirs(os.path.join(temp_dir, "images"), exist_ok=True)
# Load model
model, device = load_model()
# Process images
predictions, image_names = process_images(
args.image_dir, model, device, args.conf_threshold)
predictions, image_names = process_images(args.image_dir, model, device)
# Extract quaternions and translations
print("Converting camera parameters to COLMAP format...")
quaternions, translations = extrinsic_to_colmap_format(predictions["extrinsic"])
# Filter and prepare points
# Filter and prepare points using the original VGGT conventions
print(f"Filtering points with confidence threshold {args.conf_threshold}% and stride {args.stride}...")
points3D, image_points2D = filter_and_prepare_points(
predictions, args.conf_threshold,
predictions,
args.conf_threshold,
mask_sky=args.mask_sky,
mask_black_bg=args.mask_black_bg,
mask_white_bg=args.mask_white_bg,
stride=args.stride)
stride=args.stride,
target_dir=temp_dir,
prediction_mode=args.prediction_mode
)
# Get image dimensions
height, width = predictions["depth"].shape[1:3]
@@ -444,4 +617,4 @@ def main():
print(f"COLMAP files successfully written to {args.output_dir}")
if __name__ == "__main__":
main()
main()