app.py

import gradio as gr
from PIL import Image
import numpy as np
import torch
import requests
import os
from io import BytesIO
import traceback

# Try to import transparent-background (official InsPyReNet package)
try:
    from transparent_background import Remover
    TRANSPARENT_BG_AVAILABLE = True
    print("✅ Official transparent-background package loaded!")
except ImportError:
    TRANSPARENT_BG_AVAILABLE = False
    print("❌ transparent-background package not found. Installing...")

class OfficialInsPyReNetProcessor:
    def __init__(self):
        self.model = None
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"Using device: {self.device}")
        
    def load_model(self):
        """Load the official InsPyReNet model via transparent-background"""
        try:
            if not TRANSPARENT_BG_AVAILABLE:
                return False, "transparent-background package not installed. Please run: pip install transparent-background"
                
            print("Loading official InsPyReNet model...")
            # Initialize the official model
            # Available models: 'base', 'fast', 'ultra'
            self.model = Remover(mode='base')  # 'base' is the standard InsPyReNet
            print("✅ Official InsPyReNet model loaded successfully!")
            return True, "Model loaded successfully"
            
        except Exception as e:
            error_msg = f"Failed to load official InsPyReNet model: {str(e)}"
            print(error_msg)
            traceback.print_exc()
            return False, error_msg
    
    def remove_background(self, image):
        """Remove background using official InsPyReNet"""
        try:
            if self.model is None:
                success, msg = self.load_model()
                if not success:
                    return None, msg
            
            print(f"Processing image with size: {image.size}")
            
            # Convert PIL to the format expected by transparent-background
            # The model expects PIL Image and returns PIL Image with transparency
            result = self.model.process(image, type='rgba')
            
            print("✅ Background removed successfully!")
            return result, "Background removed successfully with official InsPyReNet!"
            
        except Exception as e:
            error_msg = f"Background removal failed: {str(e)}"
            print(error_msg)
            traceback.print_exc()
            return None, error_msg

# Alternative implementation using direct model download if transparent-background fails
class DirectInsPyReNetProcessor:
    def __init__(self):
        self.model = None
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
    def download_model(self):
        """Download official InsPyReNet checkpoint"""
        try:
            model_url = "https://github.com/plemeri/InSPyReNet/releases/download/v1.0.0/InSPyReNet_SwinB_Ultra.pth"
            model_path = "inspyrenet_model.pth"
            
            if not os.path.exists(model_path):
                print("Downloading official InsPyReNet checkpoint...")
                response = requests.get(model_url, stream=True)
                response.raise_for_status()
                
                with open(model_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)
                print("✅ Model downloaded successfully!")
            
            return model_path
            
        except Exception as e:
            print(f"Model download failed: {e}")
            return None
    
    def load_model(self):
        """Load the downloaded InsPyReNet model"""
        try:
            # This would require the full model architecture implementation
            # For now, return False to use the transparent-background package
            return False, "Please use transparent-background package for official implementation"
        except Exception as e:
            return False, str(e)

# Global processor instances
official_processor = OfficialInsPyReNetProcessor()
direct_processor = DirectInsPyReNetProcessor()

def process_image_with_inspyrenet(image):
    """Main processing function using official InsPyReNet"""
    if image is None:
        return None, "Please upload an image first."
    
    try:
        # Try official transparent-background package first
        result, message = official_processor.remove_background(image)
        
        if result is not None:
            return result, message
        else:
            return None, f"❌ {message}"
            
    except Exception as e:
        error_msg = f"Processing failed: {str(e)}"
        print(error_msg)
        traceback.print_exc()
        return None, f"❌ {error_msg}"

def install_dependencies():
    """Install required dependencies"""
    try:
        import subprocess
        import sys
        
        print("Installing transparent-background package...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "transparent-background"])
        
        # Restart may be required for import
        return "✅ Dependencies installed! Please restart the application."
    except Exception as e:
        return f"❌ Installation failed: {str(e)}"

# Custom CSS
custom_css = """
.gradio-container {
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    max-width: 1200px;
    margin: 0 auto;
}

.main-header {
    text-align: center;
    padding: 3rem 2rem;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 20px;
    color: white;
    margin-bottom: 2rem;
    box-shadow: 0 10px 40px rgba(102, 126, 234, 0.3);
}

.main-header h1 {
    font-size: 3rem;
    font-weight: 800;
    margin-bottom: 1rem;
    text-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

.main-header p {
    font-size: 1.2rem;
    opacity: 0.9;
    margin: 0.5rem 0;
    font-weight: 400;
}

.status-box {
    padding: 1rem;
    margin: 1rem 0;
    border-radius: 8px;
    font-weight: 500;
}

.status-success {
    background-color: #d1fae5;
    border: 1px solid #10b981;
    color: #047857;
}

.status-error {
    background-color: #fee2e2;
    border: 1px solid #ef4444;
    color: #dc2626;
}

.status-warning {
    background-color: #fef3c7;
    border: 1px solid #f59e0b;
    color: #92400e;
}

.feature-grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(300px, 1fr));
    gap: 1.5rem;
    margin: 2rem 0;
}

.feature-card {
    background: white;
    padding: 2rem;
    border-radius: 16px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.08);
    border: 1px solid #e5e7eb;
    text-align: center;
}

.feature-icon {
    font-size: 3rem;
    margin-bottom: 1rem;
}

.gr-button-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 16px 32px !important;
    font-weight: 600 !important;
    font-size: 1.1rem !important;
    box-shadow: 0 4px 20px rgba(102, 126, 234, 0.3) !important;
}

.gr-button-secondary {
    background: linear-gradient(135deg, #f59e0b 0%, #d97706 100%) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 12px 24px !important;
    font-weight: 600 !important;
    color: white !important;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, title="Official InsPyReNet - ACCV 2022") as demo:
    
    gr.HTML("""
        <div class="main-header">
            <h1>🏆 Official InsPyReNet</h1>
            <p><strong>Revisiting Image Pyramid Structure for High Resolution Salient Object Detection</strong></p>
            <p>ACCV 2022 • Official PyTorch Implementation</p>
            <p>Authors: Taehun Kim, Kunhee Kim, Joonyeong Lee, Dongmin Cha, Jiho Lee, Daijin Kim</p>
        </div>
    """)
    
    # Installation status
    if not TRANSPARENT_BG_AVAILABLE:
        gr.HTML(f"""
            <div class="status-box status-warning">
                <strong>⚠️ Setup Required:</strong> The official transparent-background package is not installed.<br>
                <strong>Please run:</strong> <code>pip install transparent-background</code><br>
                Or use the install button below.
            </div>
        """)
        
        with gr.Row():
            install_btn = gr.Button("📦 Install Dependencies", variant="secondary")
            install_status = gr.Textbox(label="Installation Status", interactive=False)
            
        install_btn.click(
            fn=install_dependencies,
            outputs=[install_status]
        )
    else:
        gr.HTML("""
            <div class="status-box status-success">
                ✅ <strong>Official InsPyReNet Ready!</strong> The transparent-background package is loaded.
            </div>
        """)
    
    # Main interface
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                type="pil", 
                label="📸 Upload Your Image",
                height=400
            )
            
            process_btn = gr.Button(
                "🚀 Remove Background with Official InsPyReNet",
                variant="primary",
                size="lg"
            )
            
            gr.HTML("""
                <div style="margin-top: 1rem; padding: 1rem; background: #f8fafc; border-radius: 8px; text-align: center;">
                    <strong>💡 Supported formats:</strong> JPG, PNG, JPEG<br>
                    <strong>📏 Recommended:</strong> High-resolution images for best results<br>
                    <strong>⚡ Model:</strong> Official InsPyReNet base model
                </div>
            """)
            
        with gr.Column():
            output_image = gr.Image(
                type="pil", 
                label="🎯 Official InsPyReNet Result",
                height=400
            )
            
            status_message = gr.Textbox(
                label="📊 Processing Status",
                interactive=False,
                show_label=True,
                lines=3
            )
    
    # Information sections
    gr.HTML("""
        <div class="feature-grid">
            <div class="feature-card">
                <div class="feature-icon">🏗️</div>
                <h3>Official Implementation</h3>
                <p>Direct from the ACCV 2022 paper authors. This is the exact same model used in their research, not a recreation.</p>
            </div>
            <div class="feature-card">
                <div class="feature-icon">🎯</div>
                <h3>High Resolution Focus</h3>
                <p>Specifically designed for high-resolution salient object detection with pyramid structure optimization.</p>
            </div>
            <div class="feature-card">
                <div class="feature-icon">🔬</div>
                <h3>Research Grade Quality</h3>
                <p>State-of-the-art performance on multiple SOD benchmarks including UHRSD, HRSOD, and DIS5K datasets.</p>
            </div>
        </div>
    """)
    
    gr.HTML("""
        <div style="text-align: center; margin: 2rem 0; padding: 2rem; background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%); border-radius: 16px;">
            <h3 style="color: #1f2937; margin-bottom: 1rem;">📜 About This Official Implementation</h3>
            <p style="color: #6b7280; margin-bottom: 1rem; line-height: 1.6;">
                This application uses the <strong>transparent-background</strong> package, which is the official PyPI distribution 
                of InsPyReNet created by the original research team. It provides the complete model architecture with 
                pre-trained weights exactly as used in the ACCV 2022 paper.
            </p>
            <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 1rem; margin-top: 1.5rem;">
                <div style="background: white; padding: 1rem; border-radius: 8px; border: 1px solid #e5e7eb;">
                    <strong style="color: #059669;">🏆 Official Model</strong><br>
                    <span style="color: #6b7280; font-size: 0.9rem;">Exact implementation from paper</span>
                </div>
                <div style="background: white; padding: 1rem; border-radius: 8px; border: 1px solid #e5e7eb;">
                    <strong style="color: #0891b2;">📚 ACCV 2022</strong><br>
                    <span style="color: #6b7280; font-size: 0.9rem;">Published at top-tier conference</span>
                </div>
                <div style="background: white; padding: 1rem; border-radius: 8px; border: 1px solid #e5e7eb;">
                    <strong style="color: #7c3aed;">🎯 High-Res SOD</strong><br>
                    <span style="color: #6b7280; font-size: 0.9rem;">Optimized for HR images</span>
                </div>
                <div style="background: white; padding: 1rem; border-radius: 8px; border: 1px solid #e5e7eb;">
                    <strong style="color: #dc2626;">🆓 Open Source</strong><br>
                    <span style="color: #6b7280; font-size: 0.9rem;">Free for research & commercial</span>
                </div>
            </div>
        </div>
    """)
    
    # Event handlers
    process_btn.click(
        fn=process_image_with_inspyrenet,
        inputs=[input_image],
        outputs=[output_image, status_message]
    )

# Launch configuration
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,  # Set to True if you want a public link
        debug=True
    )