131 lines
3.8 KiB
Python
131 lines
3.8 KiB
Python
"""
|
||
Local test for SmartGridCutter node logic — runs without ComfyUI.
|
||
Uses PIL (Pillow) and numpy only.
|
||
"""
|
||
import sys
|
||
sys.path.insert(0, '.')
|
||
|
||
import numpy as np
|
||
from PIL import Image
|
||
import torch
|
||
import os
|
||
|
||
# Import our node
|
||
from nodes import SmartGridCutter
|
||
|
||
def test_with_real_image():
|
||
"""Test with actual storyboard image."""
|
||
img_path = os.path.join('..', 'test_videos', 'cf_pipeline', 'storyboard.png')
|
||
|
||
if not os.path.exists(img_path):
|
||
print(f"❌ Test image not found: {img_path}")
|
||
return False
|
||
|
||
print(f"📷 Loading: {img_path}")
|
||
pil_img = Image.open(img_path).convert('RGB')
|
||
print(f" Size: {pil_img.size[0]}×{pil_img.size[1]}")
|
||
|
||
# Convert to ComfyUI tensor format [B, H, W, C]
|
||
arr = np.array(pil_img).astype(np.float32) / 255.0
|
||
tensor = torch.from_numpy(arr).unsqueeze(0)
|
||
print(f" Tensor: {tensor.shape}")
|
||
|
||
# Test auto mode
|
||
cutter = SmartGridCutter()
|
||
result = cutter.cut_grid(tensor, mode="auto")
|
||
|
||
# Unpack results
|
||
*frames, cols, rows, total, info = result
|
||
|
||
print(f"\n✅ РЕЗУЛЬТАТ:")
|
||
print(f" {info}")
|
||
print(f" Grid: {cols}×{rows} = {total} frames")
|
||
|
||
# Show frame sizes
|
||
for i in range(total):
|
||
f = frames[i]
|
||
if f.shape[1] > 16: # Not a placeholder
|
||
print(f" Frame {i}: {f.shape[2]}×{f.shape[1]}px ({f.shape})")
|
||
|
||
# Save frames for visual inspection
|
||
out_dir = os.path.join('..', 'test_videos', 'cf_pipeline', 'cut_test')
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
for i in range(total):
|
||
f = frames[i]
|
||
if f.shape[1] > 16:
|
||
f_arr = (f[0].numpy() * 255).clip(0, 255).astype(np.uint8)
|
||
f_img = Image.fromarray(f_arr)
|
||
path = os.path.join(out_dir, f'frame_{i:02d}.png')
|
||
f_img.save(path)
|
||
|
||
print(f"\n 💾 Saved {total} frames to {out_dir}/")
|
||
|
||
# Test 3x3 preset mode
|
||
print(f"\n🔧 Testing 3x3 preset mode...")
|
||
result2 = cutter.cut_grid(tensor, mode="3x3")
|
||
*frames2, cols2, rows2, total2, info2 = result2
|
||
print(f" {info2}")
|
||
|
||
return True
|
||
|
||
|
||
def test_with_synthetic():
|
||
"""Test with a synthetic grid image."""
|
||
print("🧪 Testing with synthetic 3×3 grid (900×900)...")
|
||
|
||
# Create a 3×3 grid with white dividers
|
||
cell_size = 290
|
||
divider = 10
|
||
total = cell_size * 3 + divider * 2 # 890
|
||
|
||
img = np.zeros((total, total, 3), dtype=np.uint8)
|
||
|
||
colors = [
|
||
(255, 0, 0), (0, 255, 0), (0, 0, 255),
|
||
(255, 255, 0), (255, 0, 255), (0, 255, 255),
|
||
(128, 0, 0), (0, 128, 0), (0, 0, 128),
|
||
]
|
||
|
||
for r in range(3):
|
||
for c in range(3):
|
||
y = r * (cell_size + divider)
|
||
x = c * (cell_size + divider)
|
||
idx = r * 3 + c
|
||
img[y:y+cell_size, x:x+cell_size] = colors[idx]
|
||
|
||
# White dividers
|
||
for i in range(2):
|
||
pos = (i + 1) * cell_size + i * divider
|
||
img[pos:pos+divider, :] = 255 # Horizontal
|
||
img[:, pos:pos+divider] = 255 # Vertical
|
||
|
||
tensor = torch.from_numpy(img.astype(np.float32) / 255.0).unsqueeze(0)
|
||
|
||
cutter = SmartGridCutter()
|
||
result = cutter.cut_grid(tensor, mode="auto")
|
||
*frames, cols, rows, total_frames, info = result
|
||
|
||
print(f" {info}")
|
||
assert cols == 3, f"Expected 3 cols, got {cols}"
|
||
assert rows == 3, f"Expected 3 rows, got {rows}"
|
||
print(" ✅ Synthetic test PASSED")
|
||
return True
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print("=" * 60)
|
||
print(" SmartGridCutter — Local Test Suite")
|
||
print("=" * 60)
|
||
|
||
ok1 = test_with_synthetic()
|
||
print()
|
||
ok2 = test_with_real_image()
|
||
|
||
print("\n" + "=" * 60)
|
||
if ok1 and ok2:
|
||
print(" ✅ ALL TESTS PASSED")
|
||
else:
|
||
print(" ⚠️ Some tests failed")
|
||
print("=" * 60)
|