import os import torch import numpy as np import matplotlib.pyplot as plt import gradio as gr from PIL import Image import torchvision.transforms as transforms from models.ProtoSAM import ProtoSAM, ALPNetWrapper, InputFactory, TYPE_ALPNET from models.grid_proto_fewshot import FewShotSeg from models.segment_anything.utils.transforms import ResizeLongestSide # Set environment variables for model caching os.environ['TORCH_HOME'] = "./pretrained_model" # Function to load the model def load_model(config): # Initial segmentation model alpnet = FewShotSeg( config["input_size"][0], config["reload_model_path"], config["model"] ) alpnet.cuda() base_model = ALPNetWrapper(alpnet) # ProtoSAM model sam_checkpoint = "pretrained_model/sam_vit_h.pth" model = ProtoSAM( image_size=(1024, 1024), coarse_segmentation_model=base_model, use_bbox=config["use_bbox"], use_points=config["use_points"], use_mask=config["use_mask"], debug=False, num_points_for_sam=1, use_cca=config["do_cca"], point_mode=config["point_mode"], use_sam_trans=True, coarse_pred_only=config["coarse_pred_only"], sam_pretrained_path=sam_checkpoint, use_neg_points=config["use_neg_points"], ) model = model.to(torch.device("cuda")) model.eval() return model # Function to preprocess images def preprocess_image(image, transform): if isinstance(image, np.ndarray): image_np = image else: # Convert PIL Image to numpy array image_np = np.array(image) # Convert to RGB if grayscale if len(image_np.shape) == 2: image_np = np.stack([image_np] * 3, axis=2) elif image_np.shape[2] == 1: image_np = np.concatenate([image_np] * 3, axis=2) # Apply transforms image_tensor = transform(image_np).unsqueeze(0) return image_tensor # Function to create overlay visualization def create_overlay(query_image, prediction, colormap='YlOrRd'): """ Create an overlay of the prediction on the query image """ # Convert tensors to numpy arrays for visualization if isinstance(query_image, torch.Tensor): query_image = query_image.cpu().squeeze().numpy() if isinstance(prediction, torch.Tensor): prediction = prediction.cpu().squeeze().numpy() # Normalize image for visualization query_image = (query_image - query_image.min()) / (query_image.max() - query_image.min() + 1e-8) # Ensure binary mask prediction = (prediction > 0).astype(np.float32) # Create mask overlay mask_cmap = plt.cm.get_cmap(colormap) pred_rgba = mask_cmap(prediction) pred_rgba[..., 3] = prediction * 0.7 # Set alpha channel # Create matplotlib figure for overlay fig, ax = plt.subplots(figsize=(10, 10)) # Handle grayscale vs RGB images if len(query_image.shape) == 2: ax.imshow(query_image, cmap='gray') else: if query_image.shape[0] == 3: # Channel-first format query_image = np.transpose(query_image, (1, 2, 0)) ax.imshow(query_image) ax.imshow(pred_rgba) ax.axis('off') plt.tight_layout() # Convert to PIL Image fig.canvas.draw() img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,)) plt.close(fig) return img # Model configuration config = { "input_size": [224], "reload_model_path": "path/to/your/model.pth", # Update with your model path "model": {"encoder": "resnet50", "decoder": "pspnet"}, "use_bbox": True, "use_points": True, "use_mask": True, "do_cca": True, "point_mode": "extreme", "coarse_pred_only": False, "use_neg_points": False, "base_model": TYPE_ALPNET } # Function to run inference def run_inference(query_image, support_image, support_mask, use_bbox, use_points, use_mask, use_cca, coarse_pred_only): try: # Update config based on user selections config["use_bbox"] = use_bbox config["use_points"] = use_points config["use_mask"] = use_mask config["do_cca"] = use_cca config["coarse_pred_only"] = coarse_pred_only # Check if CUDA is available if not torch.cuda.is_available(): return None, "CUDA is not available. This demo requires GPU support." # Load the model model = load_model(config) # Preprocess images sam_trans = ResizeLongestSide(1024) # Transform for images transform = transforms.Compose([ transforms.ToTensor(), transforms.Resize((1024, 1024), antialias=True), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # Process query image query_img_tensor = preprocess_image(query_image, transform) # Process support image support_img_tensor = preprocess_image(support_image, transform) # Process support mask (should be binary) support_mask_np = np.array(support_mask) support_mask_np = (support_mask_np > 127).astype(np.float32) # Binarize mask support_mask_tensor = torch.from_numpy(support_mask_np).unsqueeze(0).unsqueeze(0) support_mask_tensor = torch.nn.functional.interpolate( support_mask_tensor, size=(1024, 1024), mode='nearest' ) # Prepare model inputs support_images = [support_img_tensor.cuda()] support_masks = [support_mask_tensor.cuda()] # Create model input coarse_model_input = InputFactory.create_input( input_type=config["base_model"], query_image=query_img_tensor.cuda(), support_images=support_images, support_labels=support_masks, isval=True, val_wsize=3, original_sz=query_img_tensor.shape[-2:], img_sz=query_img_tensor.shape[-2:], gts=None, ) coarse_model_input.to(torch.device("cuda")) # Run inference with torch.no_grad(): query_pred, scores = model( query_img_tensor.cuda(), coarse_model_input, degrees_rotate=0 ) # Create overlay visualization result_image = create_overlay(query_img_tensor, query_pred) confidence_score = np.mean(scores) return result_image, f"Confidence Score: {confidence_score:.4f}" except Exception as e: return None, f"Error during inference: {str(e)}" # Define the Gradio interface def create_interface(): with gr.Blocks(title="ProtoSAM Segmentation Demo") as demo: gr.Markdown("# ProtoSAM Segmentation Demo") gr.Markdown("Upload a query image, support image, and support mask to generate a segmentation prediction.") with gr.Row(): with gr.Column(): query_image = gr.Image(label="Query Image", type="pil") support_image = gr.Image(label="Support Image", type="pil") support_mask = gr.Image(label="Support Mask", type="pil") with gr.Column(): result_image = gr.Image(label="Prediction Result") result_text = gr.Textbox(label="Result Information") with gr.Row(): with gr.Column(): use_bbox = gr.Checkbox(label="Use Bounding Box", value=True) use_points = gr.Checkbox(label="Use Points", value=True) use_mask = gr.Checkbox(label="Use Mask", value=True) with gr.Column(): use_cca = gr.Checkbox(label="Use CCA", value=True) coarse_pred_only = gr.Checkbox(label="Coarse Prediction Only", value=False) run_button = gr.Button("Run Inference") run_button.click( fn=run_inference, inputs=[ query_image, support_image, support_mask, use_bbox, use_points, use_mask, use_cca, coarse_pred_only ], outputs=[result_image, result_text] ) return demo # Create and launch the interface if __name__ == "__main__": demo = create_interface() demo.launch(share=True)