IZERE HIRWA Roger commited on
Commit
02d423d
·
1 Parent(s): 5980c9f
Files changed (3) hide show
  1. app.py +221 -0
  2. requirements.txt +11 -0
  3. sam_vit_b_01ec64.pth +3 -0
app.py ADDED
@@ -0,0 +1,221 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import cv2
3
+ import torch
4
+ from segment_anything import SamPredictor, sam_model_registry
5
+ from diffusers import StableDiffusionInpaintPipeline
6
+ import gradio as gr
7
+ from PIL import Image, ImageDraw
8
+ import matplotlib.pyplot as plt
9
+ from matplotlib.colors import LinearSegmentedColormap
10
+ import tempfile
11
+ import os
12
+
13
+ # Initialize models (cached for performance)
14
+ sam_model = None
15
+ sd_pipe = None
16
+ predictor = None
17
+
18
+ def load_models():
19
+ global sam_model, sd_pipe, predictor
20
+ if sam_model is None:
21
+ print("Loading SAM model...")
22
+ sam_model = sam_model_registry["vit_h"](checkpoint="sam_vit_b_01ec64.pth")
23
+ predictor = SamPredictor(sam_model)
24
+
25
+ if sd_pipe is None:
26
+ print("Loading Stable Diffusion model...")
27
+ sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
28
+ "stabilityai/stable-diffusion-2-inpainting",
29
+ torch_dtype=torch.float16,
30
+ safety_checker=None
31
+ ).to("cuda")
32
+
33
+ return predictor, sd_pipe
34
+
35
+ def process_house_image(image, scale):
36
+ predictor, _ = load_models()
37
+
38
+ # Convert to RGB and numpy array
39
+ image = np.array(image)
40
+ if image.shape[-1] == 4: # Remove alpha channel if exists
41
+ image = image[..., :3]
42
+
43
+ # Process with SAM
44
+ predictor.set_image(image)
45
+ masks, scores, _ = predictor.predict()
46
+
47
+ # Filter for roof segments
48
+ roof_masks = []
49
+ for i, mask in enumerate(masks):
50
+ # Basic roof filtering (position and size)
51
+ y_indices, x_indices = np.where(mask)
52
+ if len(y_indices) == 0:
53
+ continue
54
+
55
+ centroid_y = np.mean(y_indices)
56
+ height_ratio = (np.max(y_indices) - np.min(y_indices)) / image.shape[0]
57
+
58
+ # Roofs are typically in upper half and cover significant vertical space
59
+ if centroid_y < image.shape[0] * 0.6 and height_ratio > 0.1:
60
+ roof_masks.append((mask, scores[i]))
61
+
62
+ # Sort by score and select top masks
63
+ roof_masks.sort(key=lambda x: x[1], reverse=True)
64
+ final_mask = np.zeros(image.shape[:2], dtype=bool)
65
+
66
+ # Combine top 3 roof masks
67
+ for mask, _ in roof_masks[:3]:
68
+ final_mask = np.logical_or(final_mask, mask)
69
+
70
+ # Create overlay visualization
71
+ overlay = image.copy()
72
+ cmap = LinearSegmentedColormap.from_list('roof_cmap', ['#00000000', '#ff000080'])
73
+ mask_rgb = (cmap(final_mask.astype(float))[..., :3] * 255).astype(np.uint8)
74
+ overlay = (0.6 * overlay + 0.4 * mask_rgb).astype(np.uint8)
75
+
76
+ # Calculate area
77
+ roof_pixels = np.sum(final_mask)
78
+ roof_area = roof_pixels / (scale ** 2) # in square meters
79
+
80
+ return overlay, final_mask, roof_area
81
+
82
+ def calculate_sheets(roof_area, sheet_width, sheet_height, waste_factor=0.15):
83
+ sheet_area = sheet_width * sheet_height
84
+ sheets = (roof_area / sheet_area) * (1 + waste_factor)
85
+ return int(np.ceil(sheets))
86
+
87
+ def generate_new_roof(image, roof_mask, pattern_prompt, sheet_width, sheet_height):
88
+ _, sd_pipe = load_models()
89
+
90
+ # Convert to PIL format
91
+ image_pil = Image.fromarray(image)
92
+
93
+ # Convert mask to PIL format
94
+ mask_pil = Image.fromarray(roof_mask.astype(np.uint8) * 255)
95
+
96
+ # Enhance prompt with sheet dimensions
97
+ enhanced_prompt = f"{pattern_prompt}, {sheet_width:.2f}m x {sheet_height:.2f}m sheets, architectural visualization, photorealistic"
98
+
99
+ # Generate new roof
100
+ result = sd_pipe(
101
+ prompt=enhanced_prompt,
102
+ image=image_pil,
103
+ mask_image=mask_pil,
104
+ num_inference_steps=30,
105
+ guidance_scale=7.5
106
+ ).images[0]
107
+
108
+ return result
109
+
110
+ def full_process(image, scale, sheet_width, sheet_height, pattern_prompt):
111
+ # Convert image to numpy array
112
+ if isinstance(image, str):
113
+ image = np.array(Image.open(image))
114
+ else:
115
+ image = np.array(image)
116
+
117
+ # Process image to get roof mask
118
+ overlay, roof_mask, roof_area = process_house_image(image, scale)
119
+
120
+ # Calculate sheets needed
121
+ sheets_needed = calculate_sheets(roof_area, sheet_width, sheet_height)
122
+
123
+ # Generate new roof visualization
124
+ new_roof_image = generate_new_roof(image, roof_mask, pattern_prompt, sheet_width, sheet_height)
125
+
126
+ # Create result visualization
127
+ fig, ax = plt.subplots(1, 3, figsize=(18, 6))
128
+
129
+ # Original with overlay
130
+ ax[0].imshow(overlay)
131
+ ax[0].set_title("Roof Segmentation")
132
+ ax[0].axis('off')
133
+
134
+ # New roof
135
+ ax[1].imshow(new_roof_image)
136
+ ax[1].set_title("New Roof Design")
137
+ ax[1].axis('off')
138
+
139
+ # Info panel
140
+ info_text = f"Roof Area: {roof_area:.2f} m²\n\n" \
141
+ f"Sheet Size: {sheet_width} × {sheet_height} m\n\n" \
142
+ f"Sheets Needed: {sheets_needed}\n\n" \
143
+ f"Pattern: {pattern_prompt}"
144
+
145
+ ax[2].text(0.1, 0.5, info_text, fontsize=12,
146
+ bbox=dict(facecolor='white', alpha=0.8))
147
+ ax[2].axis('off')
148
+ plt.tight_layout()
149
+
150
+ # Save to temp file
151
+ temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
152
+ plt.savefig(temp_file.name, bbox_inches='tight')
153
+ plt.close()
154
+
155
+ return temp_file.name, sheets_needed
156
+
157
+ # Gradio interface
158
+ with gr.Blocks(title="Roof Renovation System", theme=gr.themes.Soft()) as demo:
159
+ gr.Markdown("# 🏠 Roof Segmentation & Renovation System")
160
+ gr.Markdown("Upload a house image, specify dimensions, and visualize new roof designs with material calculations")
161
+
162
+ with gr.Row():
163
+ with gr.Column():
164
+ image_input = gr.Image(label="Upload House Image", type="filepath")
165
+ scale_input = gr.Slider(1, 500, value=100,
166
+ label="Scale (pixels per meter)",
167
+ info="Adjust based on image perspective")
168
+
169
+ pattern_prompt = gr.Textbox(label="Roof Pattern Description",
170
+ value="modern red tile pattern",
171
+ placeholder="Describe the new roof pattern")
172
+
173
+ with gr.Row():
174
+ sheet_width = gr.Number(label="Sheet Width (meters)", value=0.5)
175
+ sheet_height = gr.Number(label="Sheet Height (meters)", value=2.0)
176
+
177
+ submit_btn = gr.Button("Generate Roof Design", variant="primary")
178
+
179
+ with gr.Column():
180
+ output_image = gr.Image(label="Results", interactive=False)
181
+ sheets_output = gr.Number(label="Sheets Needed", interactive=False)
182
+
183
+ # Examples
184
+ gr.Examples(
185
+ examples=[
186
+ ["examples/house1.jpg", 120, 0.6, 1.8, "gray shingle pattern"],
187
+ ["examples/house2.jpg", 150, 0.4, 2.2, "terracotta tile pattern"],
188
+ ["examples/house3.jpg", 200, 0.5, 2.0, "metal roof with visible seams"]
189
+ ],
190
+ inputs=[image_input, scale_input, sheet_width, sheet_height, pattern_prompt],
191
+ outputs=[output_image, sheets_output],
192
+ fn=full_process,
193
+ cache_examples=True
194
+ )
195
+
196
+ submit_btn.click(
197
+ fn=full_process,
198
+ inputs=[image_input, scale_input, sheet_width, sheet_height, pattern_prompt],
199
+ outputs=[output_image, sheets_output]
200
+ )
201
+
202
+ gr.Markdown("### How It Works:")
203
+ gr.Markdown("1. Upload a house image (aerial or perspective view) \n"
204
+ "2. Adjust the scale (pixels per meter) based on image perspective \n"
205
+ "3. Enter new roof sheet dimensions \n"
206
+ "4. Describe your desired roof pattern \n"
207
+ "5. Click generate to see the new design and material requirements")
208
+
209
+ gr.Markdown("### Technical Notes:")
210
+ gr.Markdown("- Uses Meta's Segment Anything Model (SAM) for roof detection \n"
211
+ "- Utilizes Stable Diffusion for realistic roof pattern generation \n"
212
+ "- Material calculations include 15% wastage factor \n"
213
+ "- Processing may take 20-40 seconds depending on image size")
214
+
215
+ # For Hugging Face Spaces deployment
216
+ if __name__ == "__main__":
217
+ # Create example images directory if needed
218
+ os.makedirs("examples", exist_ok=True)
219
+
220
+ # Run the app
221
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gradio>=4.0
2
+ torch
3
+ torchvision
4
+ segment-anything
5
+ diffusers
6
+ transformers
7
+ accelerate
8
+ matplotlib
9
+ numpy
10
+ opencv-python
11
+ Pillow
sam_vit_b_01ec64.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912
3
+ size 375042383