96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
import os
|
|
import sys
|
|
import numpy as np
|
|
from PIL import Image
|
|
import cv2
|
|
import torch
|
|
|
|
# Add the current directory to the path so we can import local modules
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
if current_dir not in sys.path:
|
|
sys.path.append(current_dir)
|
|
|
|
# Create a mock ProgressBar class to replace ComfyUI's ProgressBar
|
|
class MockProgressBar:
|
|
def __init__(self, total):
|
|
self.total = total
|
|
self.current = 0
|
|
|
|
def update(self, n=1):
|
|
self.current += n
|
|
print(f"Progress: {self.current}/{self.total}")
|
|
|
|
# Add the mock to sys.modules to avoid import errors
|
|
sys.modules['comfy.utils'] = type('MockComfyUtils', (), {'ProgressBar': MockProgressBar})
|
|
|
|
# Now try to import our SBS node
|
|
try:
|
|
from sbs_v2 import SBS_V2_by_SamSeen
|
|
print("Successfully imported SBS_V2_by_SamSeen")
|
|
except ImportError as e:
|
|
print(f"Error importing SBS_V2_by_SamSeen: {e}")
|
|
sys.exit(1)
|
|
|
|
def main():
|
|
"""
|
|
Test the SBS node with a sample image.
|
|
"""
|
|
# Check if an image path is provided
|
|
if len(sys.argv) > 1:
|
|
image_path = sys.argv[1]
|
|
else:
|
|
# Look for any image in the current directory
|
|
image_files = [f for f in os.listdir('.') if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp'))]
|
|
if image_files:
|
|
image_path = image_files[0]
|
|
print(f"Using image: {image_path}")
|
|
else:
|
|
print("No image files found in the current directory.")
|
|
print("Please provide an image path as an argument.")
|
|
sys.exit(1)
|
|
|
|
if not os.path.exists(image_path):
|
|
print(f"Error: Image file not found at {image_path}")
|
|
sys.exit(1)
|
|
|
|
# Load the image
|
|
img = Image.open(image_path).convert('RGB')
|
|
img_np = np.array(img)
|
|
|
|
print(f"Image shape: {img_np.shape}")
|
|
|
|
# Convert to tensor format (simulate ComfyUI input)
|
|
import torch
|
|
img_tensor = torch.tensor(img_np.astype(np.float32) / 255.0).unsqueeze(0) # Add batch dimension
|
|
|
|
# Create an instance of our SBS node
|
|
sbs_node = SBS_V2_by_SamSeen()
|
|
|
|
# Set parameters
|
|
depth_scale = 30.0
|
|
blur_radius = 15
|
|
invert_depth = False # Don't invert the depth map again
|
|
mode = "Cross-eyed"
|
|
|
|
# Process the image
|
|
print(f"Processing image with parameters: depth_scale={depth_scale}, blur_radius={blur_radius}, invert_depth={invert_depth}, mode={mode}")
|
|
sbs_image, depth_map = sbs_node.process(img_tensor, depth_scale, blur_radius, invert_depth, mode)
|
|
|
|
# Convert tensors back to numpy arrays
|
|
sbs_np = sbs_image.cpu().numpy().squeeze(0) * 255.0
|
|
depth_np = depth_map.cpu().numpy().squeeze(0) * 255.0
|
|
|
|
# Save the results
|
|
sbs_output = Image.fromarray(sbs_np.astype(np.uint8))
|
|
sbs_output.save('sbs_result.png')
|
|
print(f"SBS image saved to: sbs_result.png")
|
|
|
|
depth_output = Image.fromarray(depth_np.astype(np.uint8))
|
|
depth_output.save('depth_result_from_node.png')
|
|
print(f"Depth map saved to: depth_result_from_node.png")
|
|
|
|
print("Test completed successfully!")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|