45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
"""
|
|
Timeout wrapper for SAM2 initialization to prevent hanging
|
|
"""
|
|
|
|
import signal
|
|
import functools
|
|
from typing import Any, Callable
|
|
|
|
|
|
class TimeoutError(Exception):
|
|
pass
|
|
|
|
|
|
def timeout(seconds: int = 120):
|
|
"""Decorator to add timeout to function calls"""
|
|
def decorator(func: Callable) -> Callable:
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs) -> Any:
|
|
# Define signal handler
|
|
def timeout_handler(signum, frame):
|
|
raise TimeoutError(f"Function {func.__name__} timed out after {seconds} seconds")
|
|
|
|
# Set signal handler
|
|
old_handler = signal.signal(signal.SIGALRM, timeout_handler)
|
|
signal.alarm(seconds)
|
|
|
|
try:
|
|
result = func(*args, **kwargs)
|
|
finally:
|
|
# Restore old handler
|
|
signal.alarm(0)
|
|
signal.signal(signal.SIGALRM, old_handler)
|
|
|
|
return result
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
@timeout(120) # 2 minute timeout
|
|
def safe_init_state(predictor, video_path: str, **kwargs) -> Any:
|
|
"""Safely initialize SAM2 state with timeout"""
|
|
return predictor.init_state(
|
|
video_path=video_path,
|
|
**kwargs
|
|
) |