import gradio as gr import torch from PIL import Image import numpy as np import time import os import spaces try: from gen2seg_sd_pipeline import gen2segSDPipeline from gen2seg_mae_pipeline import gen2segMAEInstancePipeline except ImportError as e: print(f"Error importing pipeline modules: {e}") print("Please ensure gen2seg_sd_pipeline.py and gen2seg_mae_pipeline.py are in the same directory.") # Optionally, raise an error or exit if pipelines are critical at startup # raise ImportError("Could not import custom pipeline modules. Check file paths.") from e from transformers import ViTMAEForPreTraining, AutoImageProcessor # --- Configuration --- MODEL_IDS = { "SD": "reachomk/gen2seg-sd", "MAE-H": "reachomk/gen2seg-mae-h" } # Check if a GPU is available and set the device accordingly DEVICE = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {DEVICE}") # --- Global Variables for Caching Pipelines --- sd_pipe_global = None mae_pipe_global = None # --- Model Loading Functions --- def get_sd_pipeline(): """Loads and caches the gen2seg Stable Diffusion pipeline.""" global sd_pipe_global if sd_pipe_global is None: model_id_sd = MODEL_IDS["SD"] print(f"Attempting to load SD pipeline from Hugging Face Hub: {model_id_sd}") try: sd_pipe_global = gen2segSDPipeline.from_pretrained( model_id_sd, use_safetensors=True, # torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, # Optional: use float16 on GPU ).to(DEVICE) print(f"SD Pipeline loaded successfully from {model_id_sd} on {DEVICE}.") except Exception as e: print(f"Error loading SD pipeline from Hugging Face Hub ({model_id_sd}): {e}") sd_pipe_global = None # Ensure it remains None on failure # Do not raise gr.Error here; let the main function handle it. return sd_pipe_global def get_mae_pipeline(): """Loads and caches the gen2seg MAE-H pipeline.""" global mae_pipe_global if mae_pipe_global is None: model_id_mae = MODEL_IDS["MAE-H"] print(f"Loading MAE-H pipeline with model {model_id_mae} on {DEVICE}...") try: model = ViTMAEForPreTraining.from_pretrained(model_id_mae) model.to(DEVICE) model.eval() # Set to evaluation mode # Load the official MAE-H image processor # Using "facebook/vit-mae-huge" as per the original app_mae.py image_processor = AutoImageProcessor.from_pretrained("facebook/vit-mae-huge") mae_pipe_global = gen2segMAEInstancePipeline(model=model, image_processor=image_processor) # The custom MAE pipeline's model is already on the DEVICE. print(f"MAE-H Pipeline with model {model_id_mae} loaded successfully on {DEVICE}.") except Exception as e: print(f"Error loading MAE-H model or pipeline from Hugging Face Hub ({model_id_mae}): {e}") mae_pipe_global = None # Ensure it remains None on failure # Do not raise gr.Error here; let the main function handle it. return mae_pipe_global # --- Unified Prediction Function --- @spaces.GPU(duration=90) def segment_image(input_image: Image.Image, model_choice: str) -> Image.Image: """ Takes a PIL Image and model choice, performs segmentation, and returns the segmented image. """ if input_image is None: raise gr.Error("No image provided. Please upload an image.") print(f"Model selected: {model_choice}") # Ensure image is in RGB format image_rgb = input_image.convert("RGB") original_resolution = image_rgb.size # (width, height) seg_array = None try: if model_choice == "SD": pipe_sd = get_sd_pipeline() if pipe_sd is None: raise gr.Error("The SD segmentation pipeline could not be loaded. " "Please check the Space logs for more details, or try again later.") print(f"Running SD inference with image size: {image_rgb.size}") start_time = time.time() with torch.no_grad(): # The gen2segSDPipeline expects a single image or a list # The pipeline's __call__ method handles preprocessing internally seg_output = pipe_sd(image_rgb, match_input_resolution=False).prediction # Output is before resize # seg_output is expected to be a numpy array (N,H,W,1) or (N,1,H,W) or tensor # Based on gen2seg_sd_pipeline.py, if output_type="np" (default), it's [N,H,W,1] # If output_type="pt", it's [N,1,H,W] # The original app_sd.py converted tensor to numpy and squeezed. if isinstance(seg_output, torch.Tensor): seg_output = seg_output.cpu().numpy() if seg_output.ndim == 4 and seg_output.shape[0] == 1: # Batch size 1 if seg_output.shape[1] == 1: # Grayscale, (1, 1, H, W) seg_array = seg_output.squeeze(0).squeeze(0).astype(np.uint8) elif seg_output.shape[-1] == 1: # Grayscale, (1, H, W, 1) seg_array = seg_output.squeeze(0).squeeze(-1).astype(np.uint8) elif seg_output.shape[1] == 3: # RGB, (1, 3, H, W) -> (H, W, 3) seg_array = np.transpose(seg_output.squeeze(0), (1, 2, 0)).astype(np.uint8) elif seg_output.shape[-1] == 3: # RGB, (1, H, W, 3) seg_array = seg_output.squeeze(0).astype(np.uint8) else: # Fallback for unexpected shapes seg_array = seg_output.squeeze().astype(np.uint8) elif seg_output.ndim == 3: # (H, W, C) or (C, H, W) seg_array = seg_output.astype(np.uint8) elif seg_output.ndim == 2: # (H,W) seg_array = seg_output.astype(np.uint8) else: raise TypeError(f"Unexpected SD segmentation output type/shape: {type(seg_output)}, {seg_output.shape}") end_time = time.time() print(f"SD Inference completed in {end_time - start_time:.2f} seconds.") elif model_choice == "MAE-H": pipe_mae = get_mae_pipeline() if pipe_mae is None: raise gr.Error("The MAE-H segmentation pipeline could not be loaded. " "Please check the Space logs for more details, or try again later.") print(f"Running MAE-H inference with image size: {image_rgb.size}") start_time = time.time() with torch.no_grad(): # The gen2segMAEInstancePipeline expects a list of images # output_type="np" returns a NumPy array pipe_output = pipe_mae([image_rgb], output_type="np") # Prediction is (batch_size, height, width, 3) for MAE prediction_np = pipe_output.prediction[0] # Get the first (and only) image prediction end_time = time.time() print(f"MAE-H Inference completed in {end_time - start_time:.2f} seconds.") if not isinstance(prediction_np, np.ndarray): # This case should ideally not be reached if output_type="np" prediction_np = prediction_np.cpu().numpy() # Ensure it's in the expected (H, W, C) format and uint8 if prediction_np.ndim == 3 and prediction_np.shape[-1] == 3: # Expected (H, W, 3) seg_array = prediction_np.astype(np.uint8) else: # Attempt to handle other shapes if necessary, or raise error raise gr.Error(f"Unexpected MAE-H prediction shape: {prediction_np.shape}. Expected (H, W, 3).") # The MAE pipeline already does gamma correction and scaling to 0-255. # It also ensures 3 channels. else: raise gr.Error(f"Invalid model choice: {model_choice}. Please select a valid model.") if seg_array is None: raise gr.Error("Segmentation array was not generated. An unknown error occurred.") print(f"Segmentation array generated with shape: {seg_array.shape}, dtype: {seg_array.dtype}") # Convert numpy array to PIL Image # Handle grayscale or RGB based on seg_array channels if seg_array.ndim == 2: # Grayscale segmented_image_pil = Image.fromarray(seg_array, mode='L') elif seg_array.ndim == 3 and seg_array.shape[-1] == 3: # RGB segmented_image_pil = Image.fromarray(seg_array, mode='RGB') elif seg_array.ndim == 3 and seg_array.shape[-1] == 1: # Grayscale with channel dim segmented_image_pil = Image.fromarray(seg_array.squeeze(-1), mode='L') else: raise gr.Error(f"Cannot convert seg_array with shape {seg_array.shape} to PIL Image.") # Resize back to original image resolution using LANCZOS for high quality segmented_image_pil = segmented_image_pil.resize(original_resolution, Image.Resampling.LANCZOS) print(f"Segmented image processed. Output size: {segmented_image_pil.size}, mode: {segmented_image_pil.mode}") return segmented_image_pil except Exception as e: print(f"Error during segmentation with {model_choice}: {e}") # Re-raise as gr.Error for Gradio to display, if not already one if not isinstance(e, gr.Error): # It's often helpful to include the type of the original exception error_type = type(e).__name__ raise gr.Error(f"An error occurred during segmentation: {error_type} - {str(e)}") else: raise e # Re-raise if it's already a gr.Error # --- Gradio Interface --- title = "gen2seg: Generative Models Enable Generalizable Instance Segmentation Demo (SD & MAE-H)" description = f"""

Upload an image and choose a model architecture to see the instance segmentation result generated by the respective model.

BIG THANKS to Huggingface for funding our demo with their Academic GPU Grant!

Paper: https://arxiv.org/abs/2505.15263

For faster inference, please check out our GitHub to run the models locally on a GPU: https://github.com/UCDvision/gen2seg or check out our Colab demo here.

If the demo experiences issues, please open an issue on our GitHub.

If you have not already, please see our webpage at https://reachomk.github.io/gen2seg.

""" article = """ """ # Define Gradio inputs input_image_component = gr.Image(type="pil", label="Input Image") model_choice_component = gr.Dropdown( choices=list(MODEL_IDS.keys()), value="SD", # Default model label="Choose Segmentation Model Architecture" ) # Define Gradio output output_image_component = gr.Image(type="pil", label="Segmented Image") # Example images (ensure these paths are correct if you upload examples to your Space) # For example, if you create an "examples" folder in your Space repo: # example_paths = [ # os.path.join("examples", "example1.jpg"), # os.path.join("examples", "example2.png") # ] # Filter out non-existent example files to prevent errors # example_paths = [ex for ex in example_paths if os.path.exists(ex)] # Base list of example image paths/URLs base_example_images = [ "cats-on-rock-1948.jpg", "dogs.png", "000000484893.jpg", "https://reachomk.github.io/gen2seg/images/comparison/vertical/7.png", "https://reachomk.github.io/gen2seg/images/comparison/horizontal/11.png", "https://reachomk.github.io/gen2seg/images/comparison/vertical/2.jpg" ] # Generate examples for each image with both model choices model_choices_for_examples = list(MODEL_IDS.keys()) # ["SD", "MAE-H"] formatted_examples = [] for img_path_or_url in base_example_images: for model_choice in model_choices_for_examples: formatted_examples.append([img_path_or_url, model_choice]) iface = gr.Interface( fn=segment_image, inputs=[input_image_component, model_choice_component], outputs=output_image_component, title=title, description=description, article=article, examples=None, #formatted_examples if formatted_examples else None, allow_flagging="never", theme="shivi/calm_seafoam" ) if __name__ == "__main__": # Optional: Pre-load a default model on startup if desired. # This can make the first inference faster but increases startup time. # print("Attempting to pre-load default SD model on startup...") try: get_sd_pipeline() # Pre-load the default SD model print("Default SD model pre-loaded successfully or was already cached.") except Exception as e: print(f"Could not pre-load default SD model: {e}") print("Launching Gradio interface...") iface.launch()