smcs/Prompted_Segmentation_Drywall

1import numpy as np2from pathlib import Path3from PIL import Image4import torch5from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation6import gradio as gr7 8# Initialize device9device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')10 11# Patch to avoid additional_chat_templates 404 error12# We need to patch the function in the module where it is USED, not just where it's defined13print("Patching transformers to avoid additional_chat_templates 404 error...")14 15import transformers.tokenization_utils_base16import transformers.utils.hub17try:18    from huggingface_hub.errors import RemoteEntryNotFoundError19except ImportError:20    # Fallback for older versions of huggingface_hub21    from huggingface_hub.utils import EntryNotFoundError as RemoteEntryNotFoundError22 23# Capture the original function carefully to avoid recursion24# We use a unique attribute to track if we've already patched it25if not hasattr(transformers.utils.hub.list_repo_templates, "_patched"):26    _original_list_repo_templates = transformers.utils.hub.list_repo_templates27else:28    # If already patched, use the stored original29    _original_list_repo_templates = transformers.utils.hub.list_repo_templates._original30 31def patched_list_repo_templates(repo_id, *args, **kwargs):32    """Patch to catch and ignore additional_chat_templates 404 errors"""33    try:34        results = []35        # Use the captured original function36        for template in _original_list_repo_templates(repo_id, *args, **kwargs):37            results.append(template)38        return results39    except (RemoteEntryNotFoundError, Exception) as e:40        # Check if this is the additional_chat_templates error41        error_str = str(e).lower()42        if "additional_chat_templates" in error_str or "404" in error_str:43            print(f"Suppressing additional_chat_templates 404 error for {repo_id}")44            return []45        raise46 47# Mark as patched and store original48patched_list_repo_templates._patched = True49patched_list_repo_templates._original = _original_list_repo_templates50 51# Apply the patch to BOTH locations52transformers.utils.hub.list_repo_templates = patched_list_repo_templates53transformers.tokenization_utils_base.list_repo_templates = patched_list_repo_templates54print("Patch applied to transformers.tokenization_utils_base.list_repo_templates")55 56# Load processor from original model57print("Loading processor from original model...")58try:59    from transformers import CLIPTokenizer, CLIPImageProcessor60    # Load components separately61    tokenizer = CLIPTokenizer.from_pretrained("CIDAS/clipseg-rd64-refined")62    image_processor = CLIPImageProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")63    processor = CLIPSegProcessor(image_processor=image_processor, tokenizer=tokenizer)64    print("Processor loaded successfully from original model components")65except Exception as e:66    print(f"Error loading processor components: {e}")67    # Fallback: try loading processor directly (should work with patch)68    processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")69    print("Processor loaded directly with patched template check")70 71# Load models72print("Loading pretrained model...")73model_pretrained = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined").to(device)74model_pretrained.eval()75 76print("Loading fine-tuned model...")77try:78    model_trained = CLIPSegForImageSegmentation.from_pretrained("smcs/clipseg_drywall").to(device)79    model_trained.eval()80    model_trained_available = True81    print("Fine-tuned model loaded successfully from smcs/clipseg_drywall")82except Exception as e:83    print(f"Warning: Could not load fine-tuned model from smcs/clipseg_drywall: {e}")84    model_trained = None85    model_trained_available = False86 87# Define prompts88PROMPTS = {89    "segment crack": "segment crack",90    "segment taping area": "segment taping area"91}92 93# Example images94example_images = [95    ["examples/crack_1.jpg"],96    ["examples/crack_2.jpg"],97    ["examples/drywall_1.jpg"],98    ["examples/drywall_2.jpg"]99]100 101 102def overlay_mask(image, mask, alpha=0.5, color=(255, 0, 0)):103    """Overlay mask on image with transparency and colored mask"""104    if mask is None:105        return image106    107    # Ensure same size108    if mask.size != image.size:109        mask = mask.resize(image.size, Image.NEAREST)110    111    # Convert mask to numpy array112    mask_array = np.array(mask.convert('L'))113    mask_binary = (mask_array > 127).astype(np.float32)114    115    # Create colored mask116    colored_mask = np.zeros((*mask_array.shape, 3), dtype=np.uint8)117    colored_mask[:, :, 0] = color[0]  # Red channel118    colored_mask[:, :, 1] = color[1]  # Green channel119    colored_mask[:, :, 2] = color[2]  # Blue channel120    121    # Convert image to numpy array122    img_array = np.array(image.convert('RGB'))123    124    # Create overlay125    overlay = img_array.copy().astype(np.float32)126    for c in range(3):127        overlay[:, :, c] = overlay[:, :, c] * (1 - alpha * mask_binary) + colored_mask[:, :, c] * (alpha * mask_binary)128    129    overlay = overlay.astype(np.uint8)130    return Image.fromarray(overlay)131 132 133def process_image(image, prompt_option):134    """135    Process an image with both pretrained and fine-tuned models.136    137    Args:138        image: PIL Image or numpy array139        prompt_option: Selected prompt option ("segment crack" or "segment taping area")140    141    Returns:142        Tuple of (pretrained_mask, trained_mask) or error message143    """144    if image is None:145        return None, None146    147    try:148        # Convert to PIL Image if needed149        if isinstance(image, np.ndarray):150            image = Image.fromarray(image)151        elif not isinstance(image, Image.Image):152            image = Image.open(image).convert('RGB')153        else:154            image = image.convert('RGB')155        156        # Get the prompt157        prompt = PROMPTS.get(prompt_option, prompt_option)158        159        # Resize image for processing160        img_orig = image.copy()161        img = img_orig.resize((352, 352), Image.BILINEAR)162        163        # Prepare inputs164        pixel_values = processor(images=[img], return_tensors="pt")['pixel_values'].to(device)165        text_inputs = processor.tokenizer(166            prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt"167        ).to(device)168        169        # Process with pretrained model170        with torch.no_grad():171            outputs_pretrained = model_pretrained(172                pixel_values=pixel_values,173                input_ids=text_inputs['input_ids'],174                attention_mask=text_inputs['attention_mask']175            )176            logits_pretrained = outputs_pretrained.logits[0].cpu().numpy()177        178        pred_mask_pretrained = torch.sigmoid(torch.from_numpy(logits_pretrained)).numpy()179        pred_mask_pretrained = (pred_mask_pretrained > 0.5).astype(np.uint8)180        181        # Resize mask back to original image size182        pred_mask_pretrained_img = Image.fromarray(pred_mask_pretrained * 255, mode='L')183        if img_orig.size != (352, 352):184            pred_mask_pretrained_img = pred_mask_pretrained_img.resize(185                (img_orig.size[0], img_orig.size[1]), Image.NEAREST186            )187        188        # Create overlay for pretrained result (blue color)189        pred_mask_pretrained_overlay = overlay_mask(img_orig.copy(), pred_mask_pretrained_img, alpha=0.5, color=(0, 100, 255))190        191        # Process with fine-tuned model if available192        if model_trained_available and model_trained is not None:193            with torch.no_grad():194                outputs_trained = model_trained(195                    pixel_values=pixel_values,196                    input_ids=text_inputs['input_ids'],197                    attention_mask=text_inputs['attention_mask']198                )199                logits_trained = outputs_trained.logits[0].cpu().numpy()200            201            pred_mask_trained = torch.sigmoid(torch.from_numpy(logits_trained)).numpy()202            pred_mask_trained = (pred_mask_trained > 0.5).astype(np.uint8)203            204            # Resize mask back to original image size205            pred_mask_trained_img = Image.fromarray(pred_mask_trained * 255, mode='L')206            if img_orig.size != (352, 352):207                pred_mask_trained_img = pred_mask_trained_img.resize(208                    (img_orig.size[0], img_orig.size[1]), Image.NEAREST209                )210            211            # Create overlay for fine-tuned result (green color)212            pred_mask_trained_overlay = overlay_mask(img_orig.copy(), pred_mask_trained_img, alpha=0.5, color=(0, 255, 0))213        else:214            # Create a placeholder image with message215            placeholder = Image.new('RGB', img_orig.size, color=(240, 240, 240))216            pred_mask_trained_overlay = placeholder217        218        return pred_mask_pretrained_overlay, pred_mask_trained_overlay219    220    except Exception as e:221        error_msg = f"Error processing image: {str(e)}"222        print(error_msg)223        return None, None224 225 226def create_interface():227    """Create the Gradio interface"""228    229    with gr.Blocks(title="CLIPSeg Image Segmentation") as demo:230        gr.Markdown(231            """232            # CLIPSeg Image Segmentation Demo233            234            This demo compares zero-shot pretrained CLIPSeg results with fine-tuned model results.235            Select an example image or upload your own, then choose a prompt to see the segmentation results.236            """237        )238        239        with gr.Row():240            with gr.Column():241                image_input = gr.Image(242                    label="Input Image",243                    type="pil",244                    height=400245                )246                247                prompt_dropdown = gr.Dropdown(248                    choices=list(PROMPTS.keys()),249                    value=list(PROMPTS.keys())[0],250                    label="Select Prompt",251                    info="Choose the segmentation prompt"252                )253                254                submit_btn = gr.Button("Segment", variant="primary")255        256        with gr.Row():257            with gr.Column():258                pretrained_output = gr.Image(259                    label="Pretrained (Zero-shot) Result",260                    type="pil",261                    height=400262                )263            264            with gr.Column():265                trained_output = gr.Image(266                    label="Fine-tuned Result" + (" (Not Available)" if not model_trained_available else ""),267                    type="pil",268                    height=400269                )270        271        if not model_trained_available:272            gr.Markdown(273                "⚠️ **Note:** Fine-tuned model could not be loaded from `smcs/clipseg_drywall`. "274                "Only pretrained results will be shown."275            )276        277        gr.Examples(278            examples=example_images,279            inputs=image_input,280            label="Example Images"281        )282        283        # Connect the function284        submit_btn.click(285            fn=process_image,286            inputs=[image_input, prompt_dropdown],287            outputs=[pretrained_output, trained_output]288        )289        290        # Also process when example is selected291        image_input.change(292            fn=process_image,293            inputs=[image_input, prompt_dropdown],294            outputs=[pretrained_output, trained_output]295        )296        297        # Process when prompt changes298        prompt_dropdown.change(299            fn=process_image,300            inputs=[image_input, prompt_dropdown],301            outputs=[pretrained_output, trained_output]302        )303    304    return demo305 306 307if __name__ == "__main__":308    demo = create_interface()309    demo.launch(share=False)310 311