config.py

"""
Path Configuration Module
Handles all path resolution for the OmniParser system.
Works on any machine regardless of installation location.
"""

import os
from pathlib import Path
import sys

# Get the project root dynamically
def get_project_root():
    """
    Get the project root directory.
    Works by finding the parent of OmniParser directory.
    """
    # If this config is in project root, use that
    config_file = Path(__file__).resolve()
    
    # Look for OmniParser directory to determine project root
    current_dir = config_file.parent
    
    # Search up the directory tree for OmniParser or omoi-v2
    for _ in range(5):  # Search up to 5 levels
        if (current_dir / 'OmniParser').exists():
            return current_dir
        if (current_dir / '.git').exists() and 'omoi' in str(current_dir).lower():
            return current_dir
        current_dir = current_dir.parent
    
    # Fallback: assume omoi-v2 is in the parent of OmniParser
    omniparser_dir = Path(__file__).parent.parent
    if omniparser_dir.name == 'OmniParser':
        return omniparser_dir.parent
    
    # Last resort: return current working directory
    return Path.cwd()

# Get all paths
PROJECT_ROOT = get_project_root()
OMNIPARSER_DIR = PROJECT_ROOT / 'OmniParser'
WEIGHTS_DIR = OMNIPARSER_DIR / 'weights'
ICON_DETECT_DIR = WEIGHTS_DIR / 'icon_detect'
ICON_CAPTION_DIR = WEIGHTS_DIR / 'icon_caption_florence'

# Data directories
DATA_DIR = PROJECT_ROOT / 'data'
CROPPED_IMAGES_DIR = PROJECT_ROOT / 'cropped_images'
TEMP_CROP_DIR = Path('/tmp') / 'omoi_cropped_images'

# Output directories
OUTPUT_DIR = PROJECT_ROOT / 'output'

# Create directories if they don't exist
for directory in [DATA_DIR, CROPPED_IMAGES_DIR, OUTPUT_DIR, TEMP_CROP_DIR]:
    directory.mkdir(parents=True, exist_ok=True)

# ============ Model Paths ============

def get_icon_detect_model_path():
    """Get YOLOv8 icon detection model path"""
    model_path = ICON_DETECT_DIR / 'model.pt'
    if not model_path.exists():
        raise FileNotFoundError(f"Icon detect model not found at: {model_path}")
    return str(model_path)

def get_icon_caption_model_path():
    """Get Florence-2 caption model path (if it exists)"""
    model_path = ICON_CAPTION_DIR / 'model.safetensors'
    if model_path.exists():
        return str(ICON_CAPTION_DIR)
    return None

# ============ Configuration ============

def get_omniparser_config():
    """Get default OmniParser configuration"""
    return {
        'som_model_path': get_icon_detect_model_path(),
        'caption_model_name': None,  # Florence removed
        'caption_model_path': get_icon_caption_model_path(),
        'device': 'cpu',
        'BOX_TRESHOLD': 0.05,
        'save_cropped_images': True,
        'cropped_images_dir': str(TEMP_CROP_DIR)
    }

# ============ Path Resolution Functions ============

def resolve_path(path_str, base_dir=None):
    """
    Resolve a path that could be relative or absolute.
    
    Args:
        path_str: Path string (relative or absolute)
        base_dir: Base directory for relative paths (defaults to PROJECT_ROOT)
    
    Returns:
        Absolute path as string
    """
    if base_dir is None:
        base_dir = PROJECT_ROOT
    
    path = Path(path_str)
    
    # If already absolute, return as is
    if path.is_absolute():
        return str(path)
    
    # Relative path: resolve from base directory
    resolved = (Path(base_dir) / path).resolve()
    return str(resolved)

def get_screenshot_path(filename='Screenshot2.png'):
    """Get path to a screenshot file"""
    return str(DATA_DIR / filename)

def get_output_path(filename):
    """Get path to an output file"""
    return str(OUTPUT_DIR / filename)

# ============ Print Configuration for Debugging ============

if __name__ == "__main__":
    print("\n" + "="*70)
    print("OmniParser Configuration")
    print("="*70)
    print(f"\nProject Root: {PROJECT_ROOT}")
    print(f"OmniParser Dir: {OMNIPARSER_DIR}")
    print(f"Weights Dir: {WEIGHTS_DIR}")
    print(f"\n[Model Paths]")
    print(f"  Icon Detect: {get_icon_detect_model_path()}")
    print(f"  Caption Model: {get_icon_caption_model_path()}")
    print(f"\n[Data Directories]")
    print(f"  Cropped Images: {CROPPED_IMAGES_DIR}")
    print(f"  Temp Crop: {TEMP_CROP_DIR}")
    print(f"  Output: {OUTPUT_DIR}")
    print(f"\n[Configuration]")
    config = get_omniparser_config()
    for key, value in config.items():
        print(f"  {key}: {value}")
    print("\n" + "="*70 + "\n")