Upload 3 files

README.md

@@ -1,13 +1,12 @@
 ---
-title: FloorAI
-emoji: 📚
-colorFrom: gray
-colorTo: pink
+title: floorplan generation
+emoji: 📏
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 4.31.4
+sdk_version: 3.12.0
+python_version: 3.9.13
 app_file: app.py
-pinned: false
+pinned: true
 license: openrail
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---

app.py

@@ -0,0 +1,117 @@
+import gradio
+import cv2
+from PIL import Image
+import numpy as np
+
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from diffusers.utils import load_image
+import torch
+import accelerate
+import transformers
+from random import randrange
+
+
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+base_model_id = "runwayml/stable-diffusion-v1-5"
+model_id = "LuyangZ/controlnet_Neufert4_64_100"
+
+controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float32)
+controlnet.to(device)
+torch.cuda.empty_cache()
+
+pipeline = StableDiffusionControlNetPipeline.from_pretrained(base_model_id , controlnet=controlnet, torch_dtype=torch.float32)
+pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
+
+pipeline = pipeline.to(device)
+torch.cuda.empty_cache()
+seed = randrange(500)
+generator = torch.Generator(device=device).manual_seed(seed)
+
+
+
+def expand2square(ol_img, background_color):
+    width, height = ol_img.size
+    
+    if width == height:
+        pad = int(width*0.2)
+        width_new = width + pad
+        halfpad = int(pad/2)
+        
+        ol_result = Image.new(ol_img.mode, (width_new, width_new), background_color)
+        ol_result.paste(ol_img, (halfpad, halfpad))
+
+        return ol_img
+    
+    elif width > height:
+        
+        pad = int(width*0.2)
+        width_new = width + pad
+        halfpad = int(pad/2)
+
+        ol_result = Image.new(ol_img.mode, (width_new, width_new), background_color)
+        ol_result.paste(ol_img, (halfpad, (width_new - height) // 2))
+
+        return ol_result
+    
+    else:
+        pad = int(height*0.2)
+        height_new = height + pad
+        halfpad = int(pad/2)
+
+        ol_result = Image.new(ol_img.mode, (height_new, height_new), background_color)
+        ol_result.paste(ol_img, ((height_new - width) // 2, halfpad))
+
+        return ol_result
+
+def clean_img(image, mask):
+    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+    mask = cv2.threshold(mask, 250, 255, cv2.THRESH_BINARY_INV)[1]
+    
+    image[mask<250]=(255,255,255)
+    image = Image.fromarray(image).convert('RGB')
+    return image
+    
+def floorplan_generation(outline, num_of_rooms):
+    new_width = 512
+    new_height = 512
+
+    outline = cv2.cvtColor(outline, cv2.COLOR_RGB2BGR)
+    outline_original = outline.copy()
+
+    gray = cv2.cvtColor(outline, cv2.COLOR_BGR2GRAY)
+    thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)[1]
+
+    x,y,w,h = cv2.boundingRect(thresh)
+    n_outline = outline_original[y:y+h, x:x+w]
+    n_outline = cv2.cvtColor(n_outline, cv2.COLOR_BGR2RGB)
+    n_outline = Image.fromarray(n_outline).convert('RGB')
+    n_outline = expand2square(n_outline, (255, 255, 255))
+    n_outline = n_outline.resize((new_width, new_height))
+
+    num_of_rooms = str(num_of_rooms)
+    validation_prompt = "floor plan," + num_of_rooms + " bedrooms"
+    validation_image = n_outline
+
+    image = pipeline(validation_prompt, 
+                     validation_image, 
+                     num_inference_steps=20, 
+                     generator=generator).images[0]
+
+    image = np.array(image)
+    mask = np.array(n_outline)
+    mask = cv2.cvtColor(mask, cv2.COLOR_RGB2BGR)
+    image = clean_img(image, mask)
+    return image
+
+
+gradio_interface = gradio.Interface(
+  fn=floorplan_generation,
+  inputs=[gradio.Image(label="Floor Plan Outline, Entrance"),
+          gradio.Textbox(type="text", label="number of rooms", placeholder="number of rooms", show_legend=True)],
+  outputs=gradio.Image(label="Generated Floor Plan"),
+  title="floorplan generation")
+
+gradio_interface.launch(enable_queue=False)

requirements.txt

@@ -0,0 +1,6 @@
+numpy==1.25.0
+opencv-python==4.8.0.74
+diffusers == 0.24.0
+torch==2.0.1
+accelerate == 0.25.0
+transformers