import numpy as np import cv2 import torch from segment_anything import SamPredictor, sam_model_registry from diffusers import StableDiffusionInpaintPipeline import gradio as gr from PIL import Image, ImageDraw import matplotlib.pyplot as plt from matplotlib.colors import LinearSegmentedColormap import tempfile import os # Initialize models (cached for performance) sam_model = None sd_pipe = None predictor = None def load_models(): global sam_model, sd_pipe, predictor if sam_model is None: print("Loading SAM model...") sam_model = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth") # Use available checkpoint missing_keys, unexpected_keys = sam_model.load_state_dict( torch.load("sam_vit_b_01ec64.pth"), strict=False ) if missing_keys or unexpected_keys: print(f"Warning: Missing keys: {missing_keys}, Unexpected keys: {unexpected_keys}") predictor = SamPredictor(sam_model) if sd_pipe is None: print("Loading Stable Diffusion model...") device = "cuda" if torch.cuda.is_available() else "cpu" sd_pipe = StableDiffusionInpaintPipeline.from_pretrained( "stabilityai/stable-diffusion-2-inpainting", torch_dtype=torch.float16 if device == "cuda" else torch.float32, safety_checker=None ).to(device) return predictor, sd_pipe def process_house_image(image, scale): predictor, _ = load_models() # Convert to RGB and numpy array image = np.array(image) if image.shape[-1] == 4: # Remove alpha channel if exists image = image[..., :3] # Process with SAM predictor.set_image(image) masks, scores, _ = predictor.predict() # Filter for roof segments roof_masks = [] for i, mask in enumerate(masks): # Basic roof filtering (position and size) y_indices, x_indices = np.where(mask) if len(y_indices) == 0: continue centroid_y = np.mean(y_indices) height_ratio = (np.max(y_indices) - np.min(y_indices)) / image.shape[0] # Roofs are typically in upper half and cover significant vertical space if centroid_y < image.shape[0] * 0.6 and height_ratio > 0.1: roof_masks.append((mask, scores[i])) # Sort by score and select top masks roof_masks.sort(key=lambda x: x[1], reverse=True) final_mask = np.zeros(image.shape[:2], dtype=bool) # Combine top 3 roof masks for mask, _ in roof_masks[:3]: final_mask = np.logical_or(final_mask, mask) # Create overlay visualization overlay = image.copy() cmap = LinearSegmentedColormap.from_list('roof_cmap', ['#00000000', '#ff000080']) mask_rgb = (cmap(final_mask.astype(float))[..., :3] * 255).astype(np.uint8) overlay = (0.6 * overlay + 0.4 * mask_rgb).astype(np.uint8) # Calculate area roof_pixels = np.sum(final_mask) roof_area = roof_pixels / (scale ** 2) # in square meters return overlay, final_mask, roof_area def calculate_sheets(roof_area, sheet_width, sheet_height, waste_factor=0.15): sheet_area = sheet_width * sheet_height sheets = (roof_area / sheet_area) * (1 + waste_factor) return int(np.ceil(sheets)) def generate_new_roof(image, roof_mask, pattern_prompt, sheet_width, sheet_height): _, sd_pipe = load_models() # Convert to PIL format image_pil = Image.fromarray(image) # Convert mask to PIL format mask_pil = Image.fromarray(roof_mask.astype(np.uint8) * 255) # Enhance prompt with sheet dimensions enhanced_prompt = f"{pattern_prompt}, {sheet_width:.2f}m x {sheet_height:.2f}m sheets, architectural visualization, photorealistic" # Generate new roof result = sd_pipe( prompt=enhanced_prompt, image=image_pil, mask_image=mask_pil, num_inference_steps=30, guidance_scale=7.5 ).images[0] return result def full_process(image, scale, sheet_width, sheet_height, pattern_prompt): # Convert image to numpy array if isinstance(image, str): image = np.array(Image.open(image)) else: image = np.array(image) # Process image to get roof mask overlay, roof_mask, roof_area = process_house_image(image, scale) # Calculate sheets needed sheets_needed = calculate_sheets(roof_area, sheet_width, sheet_height) # Generate new roof visualization new_roof_image = generate_new_roof(image, roof_mask, pattern_prompt, sheet_width, sheet_height) # Create result visualization fig, ax = plt.subplots(1, 3, figsize=(18, 6)) # Original with overlay ax[0].imshow(overlay) ax[0].set_title("Roof Segmentation") ax[0].axis('off') # New roof ax[1].imshow(new_roof_image) ax[1].set_title("New Roof Design") ax[1].axis('off') # Info panel info_text = f"Roof Area: {roof_area:.2f} m²\n\n" \ f"Sheet Size: {sheet_width} × {sheet_height} m\n\n" \ f"Sheets Needed: {sheets_needed}\n\n" \ f"Pattern: {pattern_prompt}" ax[2].text(0.1, 0.5, info_text, fontsize=12, bbox=dict(facecolor='white', alpha=0.8)) ax[2].axis('off') plt.tight_layout() # Save to temp file temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False) plt.savefig(temp_file.name, bbox_inches='tight') plt.close() return temp_file.name, sheets_needed # Gradio interface with gr.Blocks(title="Roof Renovation System", theme=gr.themes.Soft()) as demo: gr.Markdown("# 🏠 Roof Segmentation & Renovation System") gr.Markdown("Upload a house image, specify dimensions, and visualize new roof designs with material calculations") with gr.Row(): with gr.Column(): image_input = gr.Image(label="Upload House Image", type="filepath") scale_input = gr.Slider(1, 500, value=100, label="Scale (pixels per meter)", info="Adjust based on image perspective") pattern_prompt = gr.Textbox(label="Roof Pattern Description", value="modern red tile pattern", placeholder="Describe the new roof pattern") with gr.Row(): sheet_width = gr.Number(label="Sheet Width (meters)", value=0.5) sheet_height = gr.Number(label="Sheet Height (meters)", value=2.0) submit_btn = gr.Button("Generate Roof Design", variant="primary") with gr.Column(): output_image = gr.Image(label="Results", interactive=False) sheets_output = gr.Number(label="Sheets Needed", interactive=False) submit_btn.click( fn=full_process, inputs=[image_input, scale_input, sheet_width, sheet_height, pattern_prompt], outputs=[output_image, sheets_output] ) gr.Markdown("### How It Works:") gr.Markdown("1. Upload a house image (aerial or perspective view) \n" "2. Adjust the scale (pixels per meter) based on image perspective \n" "3. Enter new roof sheet dimensions \n" "4. Describe your desired roof pattern \n" "5. Click generate to see the new design and material requirements") gr.Markdown("### Technical Notes:") gr.Markdown("- Uses Meta's Segment Anything Model (SAM) for roof detection \n" "- Utilizes Stable Diffusion for realistic roof pattern generation \n" "- Material calculations include 15% wastage factor \n" "- Processing may take 20-40 seconds depending on image size") # For Hugging Face Spaces deployment if __name__ == "__main__": # Create example images directory if needed os.makedirs("examples", exist_ok=True) # Run the app demo.launch()