64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Test script to verify installation and GPU setup"""
|
|
|
|
import sys
|
|
import torch
|
|
import cv2
|
|
import numpy as np
|
|
|
|
print("VR180 Matting Installation Test")
|
|
print("=" * 50)
|
|
|
|
# Check Python version
|
|
print(f"Python version: {sys.version}")
|
|
|
|
# Check PyTorch and CUDA
|
|
print(f"\nPyTorch version: {torch.__version__}")
|
|
print(f"CUDA available: {torch.cuda.is_available()}")
|
|
if torch.cuda.is_available():
|
|
print(f"CUDA version: {torch.version.cuda}")
|
|
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
|
print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
|
|
|
|
# Check OpenCV
|
|
print(f"\nOpenCV version: {cv2.__version__}")
|
|
|
|
# Test imports
|
|
try:
|
|
from ultralytics import YOLO
|
|
print("\n✅ YOLO import successful")
|
|
except ImportError as e:
|
|
print(f"\n❌ YOLO import failed: {e}")
|
|
|
|
try:
|
|
from sam2.build_sam import build_sam2_video_predictor
|
|
print("✅ SAM2 import successful")
|
|
except ImportError as e:
|
|
print(f"❌ SAM2 import failed: {e}")
|
|
print(" Install with: pip install git+https://github.com/facebookresearch/segment-anything-2.git")
|
|
|
|
try:
|
|
from vr180_matting.config import VR180Config
|
|
from vr180_matting.detector import YOLODetector
|
|
from vr180_matting.sam2_wrapper import SAM2VideoMatting
|
|
from vr180_matting.memory_manager import VRAMManager
|
|
print("✅ VR180 matting modules import successful")
|
|
except ImportError as e:
|
|
print(f"❌ VR180 matting import failed: {e}")
|
|
print(" Make sure to run: pip install -e .")
|
|
|
|
# Memory check
|
|
if torch.cuda.is_available():
|
|
print(f"\n📊 Current GPU Memory Usage:")
|
|
print(f" Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
|
|
print(f" Reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
|
|
print(f" Free: {(torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_reserved()) / 1024**3:.2f} GB")
|
|
|
|
print("\n" + "=" * 50)
|
|
print("Installation test complete!")
|
|
print("\nNext steps:")
|
|
print("1. If any imports failed, install missing dependencies")
|
|
print("2. Download SAM2 model weights if needed")
|
|
print("3. Run: vr180-matting --generate-config config.yaml")
|
|
print("4. Edit config.yaml with your video path")
|
|
print("5. Run: vr180-matting config.yaml") |