Upload code/cube3d/config.py

code/cube3d/config.py

@@ -0,0 +1,351 @@
+"""
+Centralized Configuration for Path Management
+
+This module provides environment-aware path management to ensure
+compatibility between local development and HuggingFace Space deployment.
+
+Usage:
+    from code.cube3d.config import DATA_DIR, LABEL_MAPPINGS, get_mapping_paths
+
+    # Get paths for a specific mapping set
+    forward_path, inverse_path = get_mapping_paths("subset_1k")
+"""
+
+import os
+import json
+import re
+import tempfile
+from pathlib import Path
+from typing import Dict, Tuple, Optional
+
+
+# ============================================================================
+# Environment Detection
+# ============================================================================
+
+def detect_environment() -> str:
+    """
+    Detect current runtime environment
+
+    Returns:
+        "huggingface" if running on HF Space, "local" otherwise
+    """
+    if os.getenv("SPACE_ID") or os.getenv("SPACE_AUTHOR_NAME"):
+        return "huggingface"
+    return "local"
+
+
+# ============================================================================
+# Path Configuration
+# ============================================================================
+
+ENVIRONMENT = detect_environment()
+
+# Project root detection
+if ENVIRONMENT == "huggingface":
+    # HuggingFace Space: app runs from /home/user/app
+    PROJECT_ROOT = Path("/home/user/app")
+else:
+    # Local: calculate from this file's location
+    # config.py is at: code/cube3d/config.py
+    # So PROJECT_ROOT = ../../.. from here
+    PROJECT_ROOT = Path(__file__).parent.parent.parent.resolve()
+
+# Data directory
+DATA_DIR = PROJECT_ROOT / "data"
+
+# Subdirectories
+CAR_1K_DIR = DATA_DIR / "car_1k"
+CAR_DATA_DIR = DATA_DIR / "1313个筛选车结构和对照渲染图"
+
+# HuggingFace model cache directory
+# Note: preload_from_hub ALWAYS saves to ~/.cache/huggingface/hub
+# It does NOT respect HF_HOME environment variable (HF official docs confirmed)
+if ENVIRONMENT == "huggingface":
+    # Use default HF cache location where preload_from_hub saves models
+    HF_CACHE_DIR = os.path.expanduser("~/.cache/huggingface")
+    print(f"[Config] Using HF cache directory: {HF_CACHE_DIR}")
+
+    # Ensure directory exists with proper permissions
+    cache_path = Path(HF_CACHE_DIR)
+    try:
+        cache_path.mkdir(parents=True, exist_ok=True, mode=0o755)
+        print(f"[Config] Cache directory created/verified with mode 755")
+    except Exception as e:
+        print(f"[Config] Warning: Could not set cache directory permissions: {e}")
+
+    # Clean stale lock files that may prevent model loading
+    try:
+        lock_files = list(cache_path.glob("**/*.lock"))
+        if lock_files:
+            print(f"[Config] Found {len(lock_files)} lock files, attempting cleanup...")
+            cleaned = 0
+            for lock_file in lock_files:
+                try:
+                    lock_file.unlink()
+                    cleaned += 1
+                except Exception as e:
+                    print(f"[Config] Could not remove {lock_file.name}: {e}")
+            print(f"[Config] Cleaned {cleaned}/{len(lock_files)} lock files")
+        else:
+            print(f"[Config] No stale lock files found")
+    except Exception as e:
+        print(f"[Config] Lock file cleanup failed: {e}")
+else:
+    # Local development
+    HF_CACHE_DIR = os.path.expanduser("~/.cache/huggingface")
+
+# Ensure cache directory exists (fallback for local)
+if ENVIRONMENT != "huggingface":
+    os.makedirs(HF_CACHE_DIR, exist_ok=True)
+
+
+# ============================================================================
+# Label Mapping Paths
+# ============================================================================
+
+LABEL_MAPPINGS: Dict[str, Dict[str, Path]] = {
+    "subset_self": {
+        "forward": CAR_1K_DIR / "subset_self" / "label_mapping.json",
+        "inverse": CAR_1K_DIR / "subset_self" / "label_inverse_mapping.json",
+    },
+    "subset_1k": {
+        "forward": CAR_1K_DIR / "subset_1k" / "label_mapping_merge.json",
+        "inverse": CAR_1K_DIR / "subset_1k" / "label_inverse_mapping_merge.json",
+    },
+}
+
+# Runtime-generated mapping cache (for HuggingFace Space with storage limits)
+_RUNTIME_MAPPING_CACHE: Dict[str, Tuple[str, str]] = {}
+
+
+# ============================================================================
+# Helper Functions
+# ============================================================================
+
+def generate_label_mappings_from_ldr(ldr_dir: Path, mapping_type: str = "subset_1k") -> Tuple[str, str]:
+    """
+    Generate label mappings by scanning LDR files at runtime
+
+    This is a fallback for HuggingFace Spaces where storage limits prevent
+    pre-uploading large mapping files. Mappings are cached in memory.
+
+    Args:
+        ldr_dir: Directory containing LDR files
+        mapping_type: Type of mapping to generate
+
+    Returns:
+        Tuple of (forward_mapping_path, inverse_mapping_path) in /tmp
+    """
+    print(f"🔧 Generating label mappings from LDR files in {ldr_dir}...")
+
+    # Check cache first
+    if mapping_type in _RUNTIME_MAPPING_CACHE:
+        print(f"✅ Using cached mappings for {mapping_type}")
+        return _RUNTIME_MAPPING_CACHE[mapping_type]
+
+    # Scan LDR files
+    label_mapping = {}  # part_name -> ID
+    label_inverse_mapping = {}  # ID -> part_name
+    label_counter = 0
+
+    ldr_files = list(ldr_dir.glob("**/*.ldr"))
+    print(f"📂 Found {len(ldr_files)} LDR files to process")
+
+    for ldr_file in ldr_files:
+        try:
+            with open(ldr_file, 'r', encoding='utf-8', errors='ignore') as f:
+                for line in f:
+                    if line.startswith('1'):  # Part data line
+                        parts = line.split()
+                        if len(parts) < 15:
+                            continue
+
+                        # Extract part identifier (lowercase, starting digits)
+                        filename = parts[14].lower()
+                        match = re.match(r'^\d+', filename)
+                        part_identifier = match.group() if match else filename
+
+                        if part_identifier not in label_mapping:
+                            label_mapping[part_identifier] = label_counter
+                            label_inverse_mapping[label_counter] = part_identifier
+                            label_counter += 1
+        except Exception as e:
+            print(f"⚠️  Error processing {ldr_file}: {e}")
+            continue
+
+    print(f"✅ Generated {len(label_mapping)} unique part mappings")
+
+    # Save to /tmp directory
+    tmp_dir = Path(tempfile.gettempdir()) / "lego_mappings" / mapping_type
+    tmp_dir.mkdir(parents=True, exist_ok=True)
+
+    forward_path = tmp_dir / "label_mapping_merge.json"
+    inverse_path = tmp_dir / "label_inverse_mapping_merge.json"
+
+    with open(forward_path, 'w', encoding='utf-8') as f:
+        json.dump(label_mapping, f, ensure_ascii=False, indent=2)
+
+    # Convert int keys to str keys for JSON
+    inverse_str_keys = {str(k): v for k, v in label_inverse_mapping.items()}
+    with open(inverse_path, 'w', encoding='utf-8') as f:
+        json.dump(inverse_str_keys, f, ensure_ascii=False, indent=2)
+
+    print(f"💾 Saved mappings to:")
+    print(f"   {forward_path}")
+    print(f"   {inverse_path}")
+
+    # Cache the paths
+    result = (str(forward_path), str(inverse_path))
+    _RUNTIME_MAPPING_CACHE[mapping_type] = result
+
+    return result
+
+
+def get_mapping_paths(mapping_type: str = "subset_1k") -> Tuple[str, str]:
+    """
+    Get label mapping file paths for a given mapping type
+
+    Automatically generates mappings from LDR files if not found.
+
+    Args:
+        mapping_type: Either "subset_self" or "subset_1k"
+
+    Returns:
+        Tuple of (forward_mapping_path, inverse_mapping_path) as strings
+
+    Raises:
+        ValueError: If mapping_type is invalid
+    """
+    if mapping_type not in LABEL_MAPPINGS:
+        raise ValueError(
+            f"Invalid mapping_type: {mapping_type}. "
+            f"Must be one of: {list(LABEL_MAPPINGS.keys())}"
+        )
+
+    forward_path = LABEL_MAPPINGS[mapping_type]["forward"]
+    inverse_path = LABEL_MAPPINGS[mapping_type]["inverse"]
+
+    # Check if files exist
+    if forward_path.exists() and inverse_path.exists():
+        return str(forward_path), str(inverse_path)
+
+    # Files don't exist - generate from LDR files as fallback
+    print(f"⚠️  Label mapping files not found for {mapping_type}")
+    print(f"   Missing: {forward_path}")
+    print(f"   Missing: {inverse_path}")
+    print(f"🔄 Generating label mappings from LDR files (this may take 1-2 minutes)...")
+
+    # Determine LDR directory to scan
+    if mapping_type == "subset_1k":
+        ldr_dir = CAR_DATA_DIR / "ldr"
+        if not ldr_dir.exists():
+            ldr_dir = CAR_DATA_DIR  # Try parent directory
+    else:
+        ldr_dir = CAR_1K_DIR / mapping_type
+
+    if not ldr_dir.exists():
+        raise FileNotFoundError(
+            f"Cannot generate mappings: LDR directory not found: {ldr_dir}\n"
+            f"Please ensure LDR files are available."
+        )
+
+    return generate_label_mappings_from_ldr(ldr_dir, mapping_type)
+
+
+def create_default_mappings(mapping_type: str = "subset_1k") -> Tuple[Dict, Dict]:
+    """
+    Create minimal default label mappings if files are missing
+
+    This is a fallback for development/testing. Production should have real files.
+
+    Args:
+        mapping_type: Mapping type identifier
+
+    Returns:
+        Tuple of (label_mapping, label_inverse_mapping) dictionaries
+    """
+    print(f"⚠️  WARNING: Creating default empty mappings for {mapping_type}")
+    print("   This is for fallback only. Production should have real mapping files.")
+
+    # Minimal mapping structure
+    label_mapping = {}
+    label_inverse_mapping = {}
+
+    return label_mapping, label_inverse_mapping
+
+
+def load_mappings_safe(mapping_type: str = "subset_1k") -> Tuple[Dict, Dict]:
+    """
+    Safely load label mappings with fallback
+
+    Attempts to load from files, falls back to defaults if missing.
+
+    Args:
+        mapping_type: Either "subset_self" or "subset_1k"
+
+    Returns:
+        Tuple of (label_mapping, label_inverse_mapping) dictionaries
+    """
+    try:
+        forward_path, inverse_path = get_mapping_paths(mapping_type)
+
+        with open(forward_path, 'r', encoding='utf-8') as f:
+            label_mapping = json.load(f)
+
+        with open(inverse_path, 'r', encoding='utf-8') as f:
+            label_inverse_mapping = json.load(f)
+
+        return label_mapping, label_inverse_mapping
+
+    except FileNotFoundError as e:
+        print(f"⚠️  {e}")
+        return create_default_mappings(mapping_type)
+
+
+# ============================================================================
+# Debug Information
+# ============================================================================
+
+def print_config_info():
+    """Print current configuration for debugging"""
+    print("=" * 60)
+    print("Configuration Information")
+    print("=" * 60)
+    print(f"Environment: {ENVIRONMENT}")
+    print(f"Project Root: {PROJECT_ROOT}")
+    print(f"Data Directory: {DATA_DIR}")
+    print(f"Data Dir Exists: {DATA_DIR.exists()}")
+    print("\nLabel Mapping Paths:")
+    for mapping_type, paths in LABEL_MAPPINGS.items():
+        print(f"\n  {mapping_type}:")
+        for key, path in paths.items():
+            exists = "✅" if path.exists() else "❌"
+            print(f"    {key}: {exists} {path}")
+    print("=" * 60)
+
+
+# ============================================================================
+# Module Test
+# ============================================================================
+
+if __name__ == "__main__":
+    print_config_info()
+
+    # Test loading mappings
+    print("\n\nTesting mapping load:")
+    try:
+        forward, inverse = get_mapping_paths("subset_1k")
+        print(f"✅ subset_1k paths retrieved successfully")
+        print(f"   Forward: {forward}")
+        print(f"   Inverse: {inverse}")
+    except Exception as e:
+        print(f"❌ Error: {e}")
+
+    try:
+        forward, inverse = get_mapping_paths("subset_self")
+        print(f"✅ subset_self paths retrieved successfully")
+        print(f"   Forward: {forward}")
+        print(f"   Inverse: {inverse}")
+    except Exception as e:
+        print(f"❌ Error: {e}")