app.py

import os
import subprocess
import sys
import shutil
import cv2
import gradio as gr
import torch
import numpy as np
from rembg import remove
from PIL import Image
import threading
import spaces
from glob import glob

# Set OpenGL Platform for headless rendering
os.environ["PYOPENGL_PLATFORM"] = "egl"

# --- Installation Helper ---
def install_dependencies():
    print("Checking and installing dependencies...", flush=True)
    try:
        # 1. Upgrade build tools (redundant but safe)
        subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "pip", "wheel", "setuptools", "ninja"])
        
        # 2. Check if Detectron2 is installed, if not install it
        # optimizing to avoid reinstall on container restart if determined present
        try:
             import detectron2
             print("Detectron2 already installed.", flush=True)
        except ImportError:
             print("Installing Detectron2...", flush=True)
             subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-build-isolation", "detectron2@git+https://github.com/facebookresearch/detectron2.git"])

        # 3. Check if ROMP is installed
        try:
             import romp
             print("ROMP already installed.", flush=True)
        except ImportError:
             print("Installing ROMP...", flush=True)
             subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-build-isolation", "git+https://github.com/ZaiqiangWu/ROMP.git#subdirectory=simple_romp"])
        
        # 4. Download Checkpoints
        if not os.path.exists("./rtv_ckpts"):
            print("Downloading checkpoints...", flush=True)
            subprocess.check_call(["git", "clone", "https://huggingface.co/wuzaiqiang/rtv_ckpts"])
        
        print("Dependencies installed and checkpoints ready.", flush=True)
    except Exception as e:
        print(f"Error installing dependencies: {e}", flush=True)

# Run installation on startup
install_dependencies()

# --- App Logic ---

# Ensure directories exist
os.makedirs("generated_masks", exist_ok=True)
os.makedirs("PerGarmentDatasets", exist_ok=True)
os.makedirs("checkpoints", exist_ok=True)

from VITON.viton_upperbody import FrameProcessor

frame_processor = None
current_garment_id = -1
garment_list = ["lab_03"] # Default pretrained

def get_garment_list():
    # Scan checkpoints for garments
    global garment_list
    ckpts = glob("checkpoints/*") + glob("rtv_ckpts/*")
    names = [os.path.basename(p) for p in ckpts if os.path.isdir(p) and "label" not in p]
    garment_list = list(set(names))
    return garment_list

def extract_frames_and_masks(video_path, garment_name):
    print(f"Processing video: {video_path}")
    cap = cv2.VideoCapture(video_path)
    mask_dir = f"generated_masks/{garment_name}"
    if os.path.exists(mask_dir):
        shutil.rmtree(mask_dir)
    os.makedirs(mask_dir)
    
    frame_count = 0
    # Limit frames for demo speed if needed, but RTV needs good coverage.
    # We'll take every frame but maybe limit total execution if video is too long.
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        # Save Mask
        # Convert to PIL for rembg
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_im = Image.fromarray(frame_rgb)
        
        # Remove background (get mask)
        # rembg returns RGBA, alpha channel is the mask
        from rembg import remove
        output = remove(pil_im)
        mask = np.array(output)[:, :, 3] # Extract alpha
        
        # Save mask as single channel png
        mask_path = os.path.join(mask_dir, f"{str(frame_count).zfill(5)}.png")
        cv2.imwrite(mask_path, mask)
        
        frame_count += 1
        if frame_count % 50 == 0:
            print(f"Processed {frame_count} frames...")
            
    cap.release()
    return mask_dir

@spaces.GPU(duration=300)
def train_garment(video_path, garment_name):
    if not video_path:
        return "Please upload a video first.", gr.Dropdown.update(choices=get_garment_list())
    
    clean_name = garment_name.strip().replace(" ", "_")
    if not clean_name:
        return "Please enter a valid garment name.", gr.Dropdown.update(choices=get_garment_list())
        
    yield f"Starting training processing for {clean_name}...", gr.Dropdown.update()
    
    # 1. Generate Masks
    try:
        yield "Creating segmentation masks (this may take a while)...", gr.Dropdown.update()
        mask_dir = extract_frames_and_masks(video_path, clean_name)
    except Exception as e:
        return f"Error during masking: {e}", gr.Dropdown.update()

    # 2. Generate Dataset
    yield "Generating dataset...", gr.Dropdown.update()
    try:
        cmd_dataset = [
            sys.executable, "DatasetGeneration/upperbody_dataset_generation.py",
            "--video_path", video_path,
            "--mask_dir", mask_dir,
            "--dataset_name", clean_name
        ]
        subprocess.check_call(cmd_dataset)
    except Exception as e:
         return f"Error during dataset generation: {e}", gr.Dropdown.update()

    # 3. Train Model
    yield "Training model (this will take minutes)...", gr.Dropdown.update()
    try:
        # Reduced params for demo speed
        cmd_train = [
            sys.executable, "Training/upperbody_training.py",
            "--model", "pix2pixHD_RGBA",
            "--input_nc", "6",
            "--output_nc", "4",
            "--batchSize", "4",
            "--img_size", "512",
            "--dataset_path", f"./PerGarmentDatasets/{clean_name}",
            "--name", clean_name,
            "--niter", "20",       # Reduced from 80 for demo
            "--niter_decay", "20"  # Reduced from 80 for demo
        ]
        subprocess.check_call(cmd_train)
    except Exception as e:
        return f"Error during training: {e}", gr.Dropdown.update()
        
    # Copy checkpoint to main checkpoints dir/ensure visibility
    # The script saves to ./checkpoints/{name} by default
    
    new_list = get_garment_list()
    return f"Training complete for {clean_name}! You can now select it in the Try-On tab.", gr.Dropdown.update(choices=new_list, value=clean_name)

# --- Inference Logic ---

def init_processor(garment_name):
    # global frame_processor, current_garment_id # Avoid global reliance in helper
    if garment_name is None:
        return None
        
    print(f"Loading garment: {garment_name}", flush=True)
    # Initialize 
    # Always create new for now to ensure we have the right one
    processor = FrameProcessor([garment_name], ckpt_dir='./checkpoints') 
    
    # Trigger load
    processor.switch_to_target_garment(0)
    return processor

@spaces.GPU
def process_frame(image, garment_name, enable_tryon):
    if image is None:
        return None
    
    if not enable_tryon:
        return image
    
    global frame_processor
    
    try:
        # Check if we need to load/reload
        # We need to treat frame_processor as potentially stale or None
        should_reload = False
        if frame_processor is None:
            should_reload = True
        elif frame_processor.garment_name_list[0] != garment_name:
             should_reload = True
             
        if should_reload:
             # Link Pretrained to checkpoints if needed
            if os.path.exists(f"rtv_ckpts/{garment_name}") and not os.path.exists(f"checkpoints/{garment_name}"):
                if not os.path.exists("checkpoints"): os.makedirs("checkpoints")
                if not os.path.exists(f"checkpoints/{garment_name}"):
                     # Copy or symlink
                     import shutil
                     shutil.copytree(f"rtv_ckpts/{garment_name}", f"checkpoints/{garment_name}")

            frame_processor = init_processor(garment_name)
    except Exception as e:
        print(f"Error loading model: {e}", flush=True)
        return image

    # Convert to RGB (Gradio is RGB, OpenCV is BGR)
    # RTV expects BGR usually? checking rtl_demo...
    # rtl_demo: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) at end.
    # FrameProcessor takes "raw_image", likely BGR from cv2.read()
    
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Pre-processing from rtl_demo
    # frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE) # Assuming webcam is standard landscape, no rotate needed usually
    # frame=resize_img(frame,max_height=1024)
    # frame=crop2_169(frame)
    # For web demo, let's keep it simple. Resizing is good.
    h, w, _ = img_bgr.shape
    if h > 1024:
        scale = 1024 / h
        img_bgr = cv2.resize(img_bgr, (int(w*scale), 1024))
        
    try:
        if frame_processor is None:
             print("Error: Frame processor is None inside process_frame")
             return image
             
        # Debug Image Stats
        h, w, c = img_bgr.shape
        print(f"Inference Input - Shape: {h}x{w}, Mean: {img_bgr.mean():.2f}, Max: {img_bgr.max()}", flush=True)
        
        output_bgr = frame_processor(img_bgr)
        
        if output_bgr is None:
             print("Warning: FrameProcessor returned None")
             # Draw Text on Image
             cv2.putText(img_bgr, "NO PERSON DETECTED", (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
             return cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
             
        return cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"Inference error: {e}")
        return image


    # --- Gradio UI ---

    with gr.Blocks(title="RTV: Real-Time Virtual Try-On") as app:
        gr.Markdown("# Real-Time Virtual Try-On")
        gr.Markdown("Train a custom garment model from a video, or try on existing ones.")

        with gr.Tab("Virtual Try-On"):
            tryon_enabled = gr.State(False)
            
            with gr.Row():
                with gr.Column():
                    garment_selector = gr.Dropdown(
                        label="Select Garment", 
                        choices=get_garment_list(),
                        value="lab_03" if "lab_03" in get_garment_list() else None,
                        interactive=True
                    )
                    input_webcam = gr.Image(sources=["webcam"], streaming=True, label="Live Feed", interactive=True)
                    with gr.Row():
                        start_btn = gr.Button("Start Try-On", variant="primary")
                        stop_btn = gr.Button("Stop Try-On", variant="stop")
                    with gr.Row():
                         stop_cam_btn = gr.Button("Turn Off Camera", variant="secondary")
                         
                with gr.Column():
                    output_display = gr.Image(label="Virtual Try-On Result")
            
            # Button handlers to toggle state
            start_btn.click(fn=lambda: True, inputs=None, outputs=[tryon_enabled])
            stop_btn.click(fn=lambda: False, inputs=None, outputs=[tryon_enabled])
            
            # Stop Camera Handler (Clears the input)
            stop_cam_btn.click(fn=lambda: None, inputs=None, outputs=[input_webcam])

            # Stream now listens to the state
            input_webcam.stream(
                fn=process_frame, 
                inputs=[input_webcam, garment_selector, tryon_enabled], 
                outputs=output_display,
                show_progress=False
            )
            
        with gr.Tab("Train New Garment"):
            gr.Markdown("Upload a video of the garment. The system will auto-mask frames and train a model.")
            with gr.Row():
                video_input = gr.Video(label="Upload Video (Mp4)")
                garment_name_input = gr.Textbox(label="Garment Name (e.g., 'red_shirt')", placeholder="my_custom_shirt")
                train_btn = gr.Button("Start Training")
            
            train_log = gr.Textbox(label="Training Status", interactive=False)
            
            train_btn.click(
                fn=train_garment,
                inputs=[video_input, garment_name_input],
                outputs=[train_log, garment_selector]
            )

    app.queue().launch()