setup_project.py

#!/usr/bin/env python3
"""
One-click setup script for U2NET MVTec LOCO Segmentation Project
Download complete project including model weights from HuggingFace
"""

import os
import sys

def setup_project():
    """
    Download and setup complete U2NET project with one command
    """
    
    print("🚀 Setting up U2NET MVTec LOCO Segmentation Project...")
    print("📦 This will download the complete project including model weights (169MB)")
    print()
    
    try:
        from huggingface_hub import snapshot_download
        print("✅ HuggingFace Hub found")
    except ImportError:
        print("❌ HuggingFace Hub not found. Installing...")
        os.system("pip install huggingface_hub")
        try:
            from huggingface_hub import snapshot_download
            print("✅ HuggingFace Hub installed successfully")
        except ImportError:
            print("❌ Failed to install HuggingFace Hub")
            return False
    
    # Repository information
    repo_id = "zhiqing0205/u2net-mvtec-loco-segmentation"
    local_dir = "./u2net-mvtec-loco"
    
    print(f"📥 Downloading from: {repo_id}")
    print(f"📁 Saving to: {local_dir}")
    print("⏳ This may take a few minutes...")
    print()
    
    try:
        # Download complete repository
        snapshot_download(
            repo_id=repo_id,
            local_dir=local_dir,
            local_dir_use_symlinks=False,
            resume_download=True
        )
        
        print("✅ Download completed successfully!")
        print(f"📂 Project location: {os.path.abspath(local_dir)}")
        print()
        print("🎯 Next steps:")
        print(f"   cd {local_dir}")
        print("   python mvtec_loco_fg_segmentation.py")
        print()
        print("📖 For detailed usage, see README.md or README_CN.md")
        
        return True
        
    except Exception as e:
        print(f"❌ Download failed: {e}")
        print()
        print("🔧 Alternative methods:")
        print("1. Manual download from: https://huggingface.co/zhiqing0205/u2net-mvtec-loco-segmentation")
        print("2. Use HuggingFace CLI: huggingface-cli download zhiqing0205/u2net-mvtec-loco-segmentation --local-dir ./u2net-project")
        return False

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(
        description="One-click setup for U2NET MVTec LOCO Segmentation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Download complete project
  python setup_project.py
  
  # Custom directory
  python setup_project.py --dir my_project
        """
    )
    
    parser.add_argument("--dir", type=str, default="./u2net-mvtec-loco",
                       help="Local directory to download project")
    
    args = parser.parse_args()
    
    # Update local directory
    repo_id = "zhiqing0205/u2net-mvtec-loco-segmentation"
    local_dir = args.dir
    
    success = setup_project()
    
    if success:
        print("🎉 Setup completed! Ready to use.")
        sys.exit(0)
    else:
        print("💡 Please try manual setup or check network connection.")
        sys.exit(1)