93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
import os
|
|
import sys
|
|
import torch
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
# Add ComfyUI path to sys.path to allow imports from custom_nodes
|
|
comfy_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
if comfy_path not in sys.path:
|
|
sys.path.append(comfy_path)
|
|
|
|
# Correct target directory for AnotherUtils
|
|
another_utils_path = os.path.dirname(os.path.abspath(__file__))
|
|
if another_utils_path not in sys.path:
|
|
sys.path.append(another_utils_path)
|
|
|
|
from inference.detection_service import DetectionService
|
|
from inference.segmentation_service import SAM2Service
|
|
from inference.depth_service import DepthAnythingService
|
|
|
|
def test_inference():
|
|
print("=== AnotherUtils Native Inference Verification (Manual Loader) ===")
|
|
|
|
# 1. Prepare Image
|
|
input_dir = os.path.join(comfy_path, "input")
|
|
try:
|
|
test_img_path = os.path.join(input_dir, "example.png")
|
|
if not os.path.exists(test_img_path):
|
|
files = [f for f in os.listdir(input_dir) if f.endswith((".png", ".jpg", ".webp"))]
|
|
if not files:
|
|
print("Error: No test image found in input folder.")
|
|
return
|
|
test_img_path = os.path.join(input_dir, files[0])
|
|
|
|
print(f"Testing with image: {test_img_path}")
|
|
img = Image.open(test_img_path).convert("RGB")
|
|
img_np = np.array(img)
|
|
except Exception as e:
|
|
print(f"Error loading image: {e}")
|
|
return
|
|
|
|
# 2. Test YOLO (Detection)
|
|
try:
|
|
print("\n--- Testing YOLOv8 (Detection) ---")
|
|
# Try finding a face model or person model
|
|
models_dir = os.path.join(comfy_path, "models", "ultralytics", "bbox")
|
|
model_path = os.path.join(models_dir, "face_yolov8m.pt")
|
|
if not os.path.exists(model_path):
|
|
model_path = os.path.join(comfy_path, "models", "ultralytics", "segm", "person_yolov8m-seg.pt")
|
|
|
|
if os.path.exists(model_path):
|
|
yolo = DetectionService(model_path)
|
|
results = yolo.infer(img_np)
|
|
print(f"Success! Found {len(results)} detections.")
|
|
else:
|
|
print(f"Warning: YOLO model not found. Skipping.")
|
|
except Exception as e:
|
|
print(f"FAILED: YOLO Error: {e}")
|
|
|
|
# 3. Test SAM 2
|
|
try:
|
|
print("\n--- Testing SAM 2 ---")
|
|
sam_dir = os.path.join(comfy_path, "models", "sam2")
|
|
sam_files = []
|
|
if os.path.exists(sam_dir):
|
|
sam_files = [f for f in os.listdir(sam_dir) if f.endswith((".safetensors", ".pt"))]
|
|
if sam_files:
|
|
sam_path = os.path.join(sam_dir, sam_files[0])
|
|
print(f"Loading SAM2 from: {sam_path}")
|
|
sam2_svc = SAM2Service(None, sam_path)
|
|
sam2_svc.set_image(img_np)
|
|
mask = sam2_svc.predict(bboxes=[[10, 10, 100, 100]])
|
|
print(f"Success! Produced mask of shape {mask.shape}")
|
|
else:
|
|
print(f"Warning: SAM2 model not found in {sam_dir}. Skipping.")
|
|
except Exception as e:
|
|
print(f"FAILED: SAM 2 Error: {e}")
|
|
|
|
# 4. Test DepthAnything V3
|
|
try:
|
|
print("\n--- Testing DepthAnything V3 ---")
|
|
depth = DepthAnythingService("v3-small")
|
|
d_map = depth.infer(img_np)
|
|
print(f"Success! Produced depth map of shape {d_map.shape}")
|
|
except Exception as e:
|
|
print(f"FAILED: Depth Error: {e}")
|
|
|
|
print("\n=== Verification Finished ===")
|
|
|
|
if __name__ == "__main__":
|
|
test_inference()
|