Upload 6 files

app_gradio.py

@@ -218,6 +218,60 @@ def create_full_mask(image):
     
     return Image.fromarray(mask).convert("RGB")
 
+def create_smart_mask(image):
+    """Create smart mask that preserves walls/windows but allows furniture placement in empty areas"""
+    import cv2
+    
+    # Convert PIL to numpy
+    img_array = np.array(image)
+    
+    # Convert to HSV for better color analysis
+    hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
+    
+    # Create mask for empty floor areas (typically lighter, less saturated)
+    # This targets floors, empty spaces where furniture can be placed
+    
+    # Method 1: Detect floor areas (usually uniform, lighter colors)
+    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+    
+    # Create base mask
+    mask = np.zeros(gray.shape, dtype=np.uint8)
+    
+    # Find relatively uniform areas (potential floor/empty spaces)
+    # Use adaptive threshold to find uniform regions
+    blur = cv2.GaussianBlur(gray, (15, 15), 0)
+    
+    # Create mask for areas that are not walls/windows (middle brightness range)
+    # Avoid very dark (shadows) and very bright (windows) areas
+    floor_mask = cv2.inRange(gray, 50, 200)
+    
+    # Remove small noise
+    kernel = np.ones((5,5), np.uint8)
+    floor_mask = cv2.morphologyEx(floor_mask, cv2.MORPH_OPEN, kernel)
+    floor_mask = cv2.morphologyEx(floor_mask, cv2.MORPH_CLOSE, kernel)
+    
+    # Create gradient mask - stronger in center, weaker at edges to preserve walls
+    h, w = floor_mask.shape
+    center_x, center_y = w//2, h//2
+    
+    # Create distance-based gradient
+    Y, X = np.ogrid[:h, :w]
+    dist_from_center = np.sqrt((X - center_x)**2 + (Y - center_y)**2)
+    max_dist = np.sqrt(center_x**2 + center_y**2)
+    
+    # Central areas get higher weight (more likely to be inpainted)
+    gradient = 1 - (dist_from_center / max_dist) * 0.7  # Keep some mask even at edges
+    gradient = np.clip(gradient, 0.2, 1.0)  # Minimum 20% mask everywhere
+    
+    # Combine floor detection with gradient
+    final_mask = (floor_mask.astype(np.float32) / 255.0) * gradient * 255
+    final_mask = final_mask.astype(np.uint8)
+    
+    # Apply some blur for smoother transitions
+    final_mask = cv2.GaussianBlur(final_mask, (11, 11), 0)
+    
+    return Image.fromarray(final_mask).convert("RGB")
+
 def post_process_blend(original, generated):
     """Post-process to reduce seam artifacts between original and generated areas"""
     from PIL import ImageFilter
@@ -253,7 +307,7 @@ def post_process_blend(original, generated):
     return Image.fromarray(final_array)
 
 @spaces.GPU
-def design_space(input_image, room_type, design_style, num_steps, guidance_scale, strength):
+def design_space(input_image, room_type, design_style, inpainting_mode, num_steps, guidance_scale, strength):
     """Generate space design using ZeroGPU"""
     
     global pipe, seg_processor, seg_model, mlsd_processor
@@ -297,8 +351,13 @@ def design_space(input_image, room_type, design_style, num_steps, guidance_scale
         else:
             mlsd_image = resized_image.copy()
         
-        # Create full mask to eliminate boundaries
-        mask_image = create_full_mask(resized_image)
+        # Create mask based on selected mode
+        if inpainting_mode == "smart":
+            mask_image = create_smart_mask(resized_image)
+            # Adjust prompt for furniture placement
+            prompt = f"photorealistic interior design, add {detailed_prompt.split(',')[0]} furniture and decor, preserve existing walls and windows, natural lighting, ultra-realistic, high resolution, sharp focus, interior design magazine quality, realistic textures, realistic materials"
+        else:
+            mask_image = create_full_mask(resized_image)
         
         # Generate image with original project settings
         result = pipe(
@@ -357,6 +416,15 @@ def create_interface():
                     label="Select Design Style"
                 )
                 
+                inpainting_mode = gr.Radio(
+                    choices=[
+                        ("Complete Room Redesign", "full"),
+                        ("Add Furniture Only (Preserve Walls)", "smart")
+                    ],
+                    value="smart",
+                    label="🎨 Design Mode",
+                    info="Choose how to modify your image"
+                )
                 
                 with gr.Accordion("⚙️ Advanced Settings", open=False):
                     num_steps = gr.Slider(
@@ -389,7 +457,7 @@ def create_interface():
         # Event handler
         generate_btn.click(
             fn=design_space,
-            inputs=[input_image, room_type, design_style, num_steps, guidance_scale, strength],
+            inputs=[input_image, room_type, design_style, inpainting_mode, num_steps, guidance_scale, strength],
             outputs=[output_image, result_message]
         )
         

requirements_gradio.txt

@@ -9,4 +9,5 @@ gradio
 spaces
 pillow
 numpy
-scipy
+scipy
+opencv-python