diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..778915e --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,164 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +ComfyUI Depth Estimation Node is a custom node for ComfyUI that provides depth map generation from images using Depth-Anything models (V1, V2, and V3). The node integrates with ComfyUI's workflow system and provides configurable post-processing options. **Version 1.3.4+ adds Depth-Anything V3 support with camera pose estimation (extrinsics/intrinsics).** + +## Development Commands + +### Installation +```bash +# Install dependencies +pip install -r requirements.txt + +# For DA3 camera pose estimation (optional) +pip install depth_anything_3 +``` + +### Testing the Node +The node is designed to run within ComfyUI. To test: +1. Copy/symlink this directory to `ComfyUI/custom_nodes/ComfyUIDepthEstimation` +2. Restart ComfyUI +3. Load the example workflow from `workflows/da3_camera_demo.json` (for DA3 features) + +### Publishing +The project uses ComfyUI Registry for distribution: +```bash +# Project metadata is in pyproject.toml +# Publishing is configured in publish.yaml +``` + +## Code Architecture + +### Entry Point: `__init__.py` +- Implements graceful dependency checking to avoid ComfyUI crashes +- Creates placeholder error nodes if dependencies are missing +- Registers the node with ComfyUI via `NODE_CLASS_MAPPINGS` and `NODE_DISPLAY_NAME_MAPPINGS` +- Falls back to error nodes with helpful messages if imports fail + +### Main Implementation: `depth_estimation_node.py` + +#### Model System +The node supports multiple depth estimation models defined in `DEPTH_MODELS` dictionary: +- **Depth-Anything V1**: Small, Base, Large variants +- **Depth-Anything V2**: Small, Base variants +- **Depth-Anything V3**: Small, Base variants (Apache 2.0, with camera pose estimation) + +Each model entry includes: HuggingFace path, VRAM requirements, direct download URL, model type, encoder type, and `supports_pose` flag. + +#### DA3 Camera Pose Estimation (New in v1.3.4) +The node now supports camera extrinsics and intrinsics estimation for DA3 models: +- **DA3Prediction dataclass**: Structured output with depth, confidence, extrinsics, intrinsics, raw_depth +- **DA3ModelWrapper**: Wraps depth_anything_3 library with pose extraction +- **5-tuple output**: `(depth, confidence, extrinsics, intrinsics, camera_json)` + +#### Main Node Class: `DepthEstimationNode` +**Input Parameters:** +- `image`: ComfyUI IMAGE tensor +- `model_name`: Dropdown of available models from DEPTH_MODELS +- `input_size`: Target size for inference (384-8192, default 1024) +- `blur_radius`: Gaussian blur radius (FLOAT, 0.0-100.0) +- `median_size`: Median filter size (INT, odd values only for PIL compatibility) +- `apply_auto_contrast`: Boolean for contrast enhancement +- `apply_gamma`: Boolean for gamma correction +- `force_reload`: Force model reload (optional) +- `force_cpu`: Force CPU processing (optional) +- `enable_camera_estimation`: Enable camera pose output for DA3 models (default True) +- `output_raw_depth`: Output raw metric depth instead of normalized (optional) + +**Output Types:** +```python +RETURN_TYPES = ("IMAGE", "IMAGE", "CAMERA_EXTRINSICS", "CAMERA_INTRINSICS", "STRING") +RETURN_NAMES = ("depth", "confidence", "extrinsics", "intrinsics", "camera_json") +``` + +**Main Method:** `estimate_depth()` +- Validates input size and parameters +- Loads/reloads models with intelligent caching +- For DA3 models: extracts camera extrinsics [N, 3, 4] and intrinsics [N, 3, 3] +- Serializes camera data to comprehensive JSON with depth/confidence stats +- Returns 5-tuple with all outputs + +#### Key Helper Methods +- `_serialize_camera_data()`: Converts camera tensors to structured JSON +- `_check_pose_support()`: Determines if model supports camera estimation +- `cleanup()`: Cleans up VRAM and manages model resources + +### Integration with ComfyUI +- Uses `folder_paths` for ComfyUI path management +- Uses `comfy.model_management` for device selection and memory management +- Custom types: `CAMERA_EXTRINSICS`, `CAMERA_INTRINSICS` for type safety +- Node registered in "depth" category + +## Camera JSON Output Format + +When using DA3 models with camera estimation enabled: +```json +{ + "format_version": "1.0", + "model": "Depth-Anything-V3-Base", + "timestamp": 1234567890.123, + "image_width": 1024, + "image_height": 1024, + "num_views": 1, + "extrinsics": [{ + "matrix": [[...], [...], [...]], + "format": "opencv_w2c", + "rotation": [[...], [...], [...]], + "translation": [x, y, z] + }], + "intrinsics": [{ + "matrix": [[...], [...], [...]], + "fx": 1176.38, + "fy": 1167.86, + "cx": 252.0, + "cy": 252.0 + }], + "depth_stats": {"min": 0.68, "max": 1.37, "mean": 0.97, "std": 0.17}, + "confidence_stats": {"mean": 1.03, "std": 0.07, "min": 1.0, "max": 1.42} +} +``` + +## Important Implementation Details + +### Model Loading Strategy +Three-tier fallback: +1. HuggingFace transformers pipeline (preferred) +2. Direct PyTorch model download from GitHub releases +3. Local cached models + +### VRAM Management +- Models cached in `self.depth_estimator` to avoid reloading +- Explicit cleanup with `torch.cuda.empty_cache()` and `gc.collect()` +- VRAM requirements: V1/V2 (1.5-4GB), DA3 (~2-2.5GB) + +### Image Format Conversions +- ComfyUI IMAGE tensors: BHWC format, 0-1 range +- PIL Images: RGB, 0-255 range +- Camera tensors: [N, 3, 4] for extrinsics, [N, 3, 3] for intrinsics + +## ComfyUI Node Conventions + +When modifying or extending this node: +- INPUT_TYPES must return a dict with "required" and optional "optional" keys +- Input types use ComfyUI type strings: "IMAGE", "INT", "FLOAT", "BOOLEAN", "STRING" +- Parameter constraints defined in dicts: `{"default": X, "min": Y, "max": Z, "step": W}` +- RETURN_TYPES is a tuple of ComfyUI type strings +- Custom types (CAMERA_EXTRINSICS, CAMERA_INTRINSICS) defined as string constants +- All nodes must be registered in `NODE_CLASS_MAPPINGS` in `__init__.py` + +## Known Issues and Workarounds + +1. **Protobuf errors**: Fixed by pinning to protobuf==3.20.3 +2. **Model download failures**: Node includes retry logic and fallback URLs +3. **VRAM issues**: Use smaller models or enable `force_cpu` flag +4. **DA3 not available**: Node gracefully hides DA3 models if depth_anything_3 not installed +5. **Median filter**: PIL requires odd integers > 1; validation added + +## File Locations + +- Example workflows: `workflows/da3_camera_demo.json`, `workflows/Depth_Map_Generator_V1.json` +- Icon: `icon.png` (referenced in pyproject.toml) +- Images: `images/` directory diff --git a/GEMINI.md b/GEMINI.md index a47bb64..bfced38 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -72,8 +72,6 @@ Critical dependencies include: - `model_name`: Selector for models (e.g., "Depth-Anything-V2-Small"). - `input_size`: Resolution for inference (default 1024). - `blur_radius`, `median_size`: Smoothing parameters. -- `input_size`: Resolution for inference (default 1024). - - `blur_radius`, `median_size`: Smoothing parameters. - `force_cpu`: Option to offload VRAM. ## Recent Updates (DA3 Integration) diff --git a/depth_estimation_node.py b/depth_estimation_node.py index 8dc6adc..f286a59 100644 --- a/depth_estimation_node.py +++ b/depth_estimation_node.py @@ -701,13 +701,15 @@ class DA3ModelWrapper: logger.info(f"DA3ModelWrapper initialized: {model_name}, Pose support: {self.supports_pose}, Device: {device}") def _check_pose_support(self, model_name: str) -> bool: + """Check if model supports camera pose estimation (DA3 variants only).""" model_lower = model_name.lower() + # Mono/metric variants don't support pose estimation if "mono" in model_lower or "metric" in model_lower: return False - if "v3" in model_lower or "da3" in model_lower: - if any(variant in model_lower for variant in ["small", "base", "large", "giant", "nested"]): - return True - return False + # DA3 models with size variants support pose estimation + is_da3 = "v3" in model_lower or "da3" in model_lower + size_variants = ["small", "base", "large", "giant", "nested"] + return is_da3 and any(v in model_lower for v in size_variants) def __call__(self, image: Union[Image.Image, List[Image.Image]]) -> DA3Prediction: try: @@ -761,14 +763,13 @@ class DA3ModelWrapper: proc_array = prediction.processed_images processed_tensor = torch.from_numpy(proc_array).to(self.device) - # Handle single image case + # Handle single image case - squeeze batch dimension from all tensors if not is_batch: - depth_tensor = depth_tensor.squeeze(0) - raw_depth_tensor = raw_depth_tensor.squeeze(0) - if confidence_tensor is not None: confidence_tensor = confidence_tensor.squeeze(0) - if extrinsics_tensor is not None: extrinsics_tensor = extrinsics_tensor.squeeze(0) - if intrinsics_tensor is not None: intrinsics_tensor = intrinsics_tensor.squeeze(0) - if processed_tensor is not None: processed_tensor = processed_tensor.squeeze(0) + tensors = [depth_tensor, raw_depth_tensor, confidence_tensor, + extrinsics_tensor, intrinsics_tensor, processed_tensor] + depth_tensor, raw_depth_tensor, confidence_tensor, \ + extrinsics_tensor, intrinsics_tensor, processed_tensor = \ + [t.squeeze(0) if t is not None else None for t in tensors] return DA3Prediction( depth=depth_tensor, @@ -848,8 +849,8 @@ class DepthEstimationNode: }, "optional": { "input_size": ("INT", {"default": 1024, "min": 384, "max": 8192, "step": 32}), - "blur_radius": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), - "median_size": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), + "blur_radius": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}), + "median_size": ("INT", {"default": 0, "min": 0, "max": 21, "step": 2}), # Odd values only, max 21 for PIL "apply_auto_contrast": ("BOOLEAN", {"default": False}), "apply_gamma": ("BOOLEAN", {"default": False}), "force_reload": ("BOOLEAN", {"default": False}), @@ -2410,12 +2411,117 @@ SOLUTION: """Convert PIL Image back to tensor.""" img_np = np.array(pil_img).astype(np.float32) / 255.0 tensor = torch.from_numpy(img_np).unsqueeze(0) - + if self.device is not None: tensor = tensor.to(self.device) - + return tensor + def _serialize_camera_data( + self, + extrinsics: Optional[torch.Tensor], + intrinsics: Optional[torch.Tensor], + confidence: Optional[torch.Tensor], + raw_depth: Optional[torch.Tensor], + model_name: str, + image_width: int, + image_height: int + ) -> str: + """ + Serialize camera data to comprehensive JSON string. + + Args: + extrinsics: [N, 3, 4] camera extrinsics (opencv w2c format) + intrinsics: [N, 3, 3] camera intrinsics + confidence: [N, H, W] confidence maps (optional) + raw_depth: [N, H, W] raw depth values (optional) + model_name: Name of the model used + image_width: Original image width + image_height: Original image height + + Returns: + JSON string with camera parameters and metadata + """ + import json + + data = { + "format_version": "1.0", + "model": model_name, + "timestamp": time.time(), + "image_width": image_width, + "image_height": image_height, + } + + # Serialize extrinsics + if extrinsics is not None: + ext_np = extrinsics.cpu().numpy() + num_views = ext_np.shape[0] if ext_np.ndim == 3 else 1 + data["num_views"] = num_views + + if ext_np.ndim == 3: # Batch [N, 3, 4] + data["extrinsics"] = [ + { + "matrix": ext.tolist(), + "format": "opencv_w2c", + "rotation": ext[:3, :3].tolist(), + "translation": ext[:3, 3].tolist() if ext.shape[1] > 3 else [0, 0, 0] + } + for ext in ext_np + ] + else: # Single [3, 4] + data["extrinsics"] = [{ + "matrix": ext_np.tolist(), + "format": "opencv_w2c", + "rotation": ext_np[:3, :3].tolist(), + "translation": ext_np[:3, 3].tolist() if ext_np.shape[1] > 3 else [0, 0, 0] + }] + + # Serialize intrinsics + if intrinsics is not None: + int_np = intrinsics.cpu().numpy() + + if int_np.ndim == 3: # Batch [N, 3, 3] + data["intrinsics"] = [ + { + "matrix": intr.tolist(), + "fx": float(intr[0, 0]), + "fy": float(intr[1, 1]), + "cx": float(intr[0, 2]), + "cy": float(intr[1, 2]) + } + for intr in int_np + ] + else: # Single [3, 3] + data["intrinsics"] = [{ + "matrix": int_np.tolist(), + "fx": float(int_np[0, 0]), + "fy": float(int_np[1, 1]), + "cx": float(int_np[0, 2]), + "cy": float(int_np[1, 2]) + }] + + # Add depth statistics + if raw_depth is not None: + depth_np = raw_depth.cpu().numpy() + data["depth_stats"] = { + "min": float(depth_np.min()), + "max": float(depth_np.max()), + "mean": float(depth_np.mean()), + "std": float(depth_np.std()) + } + + # Add confidence statistics if available + if confidence is not None: + conf_np = confidence.cpu().numpy() + data["confidence_stats"] = { + "mean": float(conf_np.mean()), + "std": float(conf_np.std()), + "min": float(conf_np.min()), + "max": float(conf_np.max()) + } + + return json.dumps(data, indent=2) + def estimate_depth(self, image: torch.Tensor, model_name: str, @@ -2458,7 +2564,12 @@ SOLUTION: # Initialize all return values to None or empty depth_output = None - confidence_output = torch.zeros((1, 512, 512, 3), dtype=torch.float32) # Placeholder for confidence + # Get input dimensions for dynamic placeholder sizing + if image is not None and hasattr(image, 'shape') and len(image.shape) >= 3: + img_h, img_w = image.shape[1], image.shape[2] + else: + img_h, img_w = 512, 512 # Fallback only if input is invalid + confidence_output = torch.zeros((1, img_h, img_w, 3), dtype=torch.float32, device=image.device if hasattr(image, 'device') else 'cpu') extrinsics_output = None intrinsics_output = None camera_json_output = "" @@ -2771,35 +2882,73 @@ SOLUTION: depth_map = predicted_depth_tensor.squeeze(0).cpu().numpy() if depth_result.confidence is not None: - confidence_output = depth_result.confidence.unsqueeze(0) # Add batch dim - confidence_output = torch.stack([confidence_output.squeeze()] * 3, dim=-1) # Convert to RGB for ComfyUI - + conf = depth_result.confidence + # Ensure batch dimension [N, H, W] + if conf.ndim == 2: + conf = conf.unsqueeze(0) # [H, W] -> [1, H, W] + # Convert to RGB format [N, H, W, 3] for ComfyUI IMAGE type + confidence_output = conf.unsqueeze(-1).repeat(1, 1, 1, 3) # [N, H, W, 3] + if enable_camera_estimation and depth_result.supports_pose: if depth_result.extrinsics is not None: - extrinsics_output = depth_result.extrinsics.unsqueeze(0) # Add batch dim + ext = depth_result.extrinsics + # Ensure batch dimension for single image case + if ext.ndim == 2: # [3, 4] -> [1, 3, 4] + extrinsics_output = ext.unsqueeze(0) + else: # Already [N, 3, 4] + extrinsics_output = ext if depth_result.intrinsics is not None: - intrinsics_output = depth_result.intrinsics.unsqueeze(0) # Add batch dim + intr = depth_result.intrinsics + # Ensure batch dimension for single image case + if intr.ndim == 2: # [3, 3] -> [1, 3, 3] + intrinsics_output = intr.unsqueeze(0) + else: # Already [N, 3, 3] + intrinsics_output = intr # Prepare camera_json_output - if extrinsics_output is not None and intrinsics_output is not None: - import json - camera_params = { - "extrinsics": extrinsics_output.squeeze(0).cpu().numpy().tolist(), - "intrinsics": intrinsics_output.squeeze(0).cpu().numpy().tolist(), - "image_width": original_width, - "image_height": original_height, - "depth_min": float(depth_result.raw_depth.min().cpu().numpy()) if depth_result.raw_depth is not None else 0.0, - "depth_max": float(depth_result.raw_depth.max().cpu().numpy()) if depth_result.raw_depth is not None else 1.0, - "model_name": model_name - } - camera_json_output = json.dumps(camera_params, indent=2) - logger.info("Camera parameters extracted and formatted to JSON.") - else: - logger.warning("Camera estimation enabled but extrinsics or intrinsics were None.") + # Prepare camera_json_output using the helper method + try: + camera_json_output = self._serialize_camera_data( + extrinsics=depth_result.extrinsics, + intrinsics=depth_result.intrinsics, + confidence=depth_result.confidence, + raw_depth=depth_result.raw_depth, + model_name=model_name, + image_width=original_width, + image_height=original_height + ) + logger.info(f"Camera JSON generated successfully. Length: {len(camera_json_output)}") + # Explicit print for user debugging as requested + print(f"DA3 Camera JSON Output (first 200 chars): {camera_json_output[:200]}...") + except Exception as json_error: + logger.error(f"Error serializing camera data: {str(json_error)}") + logger.error(traceback.format_exc()) + camera_json_output = json.dumps({"error": str(json_error)}) else: logger.info("Camera estimation disabled or not supported by model.") + # Still output informative JSON even without camera data + camera_json_output = json.dumps({ + "status": "camera_estimation_unavailable", + "reason": "disabled" if not enable_camera_estimation else "model_does_not_support_pose", + "model": model_name, + "supports_pose": depth_result.supports_pose, + "enable_camera_estimation": enable_camera_estimation, + "extrinsics_available": depth_result.extrinsics is not None, + "intrinsics_available": depth_result.intrinsics is not None, + "image_width": original_width, + "image_height": original_height + }, indent=2) else: # Handle V1/V2/MiDaS pipeline output + # V1/V2/MiDaS don't support camera estimation - output informative JSON + camera_json_output = json.dumps({ + "status": "not_supported", + "reason": "model_type_does_not_support_camera_estimation", + "model": model_name, + "message": "Camera pose estimation requires Depth-Anything-V3 models", + "image_width": original_width, + "image_height": original_height + }, indent=2) # Verify depth result and convert to float32 if not isinstance(depth_result, dict) or "predicted_depth" not in depth_result: logger.error(f"Invalid depth result format: {type(depth_result)}") @@ -2927,10 +3076,13 @@ SOLUTION: except Exception as blur_error: logger.warning(f"Error applying blur: {str(blur_error)}. Skipping.") - # Apply median filter if size is valid + # Apply median filter if size is valid (must be odd number > 1) try: - median_size_int = int(median_size_str) - if median_size_int > 0: + median_size_int = int(median_size) if median_size is not None else 0 + # Ensure median size is odd (PIL requirement) - round up to nearest odd + if median_size_int > 1: + if median_size_int % 2 == 0: + median_size_int += 1 # Make odd depth_pil = depth_pil.filter(ImageFilter.MedianFilter(size=median_size_int)) except Exception as median_error: logger.warning(f"Error applying median filter: {str(median_error)}. Skipping.") @@ -3006,7 +3158,7 @@ SOLUTION: processing_time = time.time() - start_time logger.info(f"Depth processing completed in {processing_time:.2f} seconds") logger.info(f"Output tensor: shape={depth_tensor.shape}, dtype={depth_tensor.dtype}, device={depth_tensor.device}") - + return (depth_tensor, confidence_output, extrinsics_output, intrinsics_output, camera_json_output) except Exception as post_error: