Update app.py

app.py

@@ -1,89 +1,93 @@
-import gradio as gr
-import torch
-import numpy as np
-import cv2
-from PIL import Image
-from torchvision import transforms
-from transformers import SegformerForSemanticSegmentation, AutoImageProcessor
-
-# Device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load model and processor once
-processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b2-finetuned-ade-512-512")
-model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b2-finetuned-ade-512-512").to(device)
-
-def process(room_img, tile_img):
-    room_img = room_img.convert("RGB")
-    tile_img = tile_img.convert("RGB")
-    room_np = np.array(room_img)
-
-    # Segmentation
-    inputs = processor(images=room_img, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    segmentation = outputs.logits.argmax(dim=1).squeeze().cpu().numpy()
-    segmentation_resized = cv2.resize(segmentation.astype(np.uint8), (room_np.shape[1], room_np.shape[0]), interpolation=cv2.INTER_NEAREST)
-
-    # Mask for floor (ADE20K class index 3)
-    floor_class_index = 3
-    mask_bin = (segmentation_resized == floor_class_index).astype(np.uint8)
-
-    # Largest contour
-    contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if not contours:
-        return room_img, Image.fromarray(mask_bin*255), tile_img, room_img
-    contour = max(contours, key=cv2.contourArea)
-    if len(contour) < 4:
-        return room_img, Image.fromarray(mask_bin*255), tile_img, room_img
-
-    x, y, w, h = cv2.boundingRect(contour)
-    src_pts = np.array([[x, y + h], [x + w, y + h], [x + w, y], [x, y]], dtype=np.float32)
-    offset = h * 0.5
-    dst_pts = np.array([[x - offset, y + h], [x + w + offset, y + h], [x + w, y], [x, y]], dtype=np.float32)
-    H = cv2.getPerspectiveTransform(src_pts, dst_pts)
-
-    # Resize tile
-    target_tile_width = room_np.shape[1] // 10
-    tile_aspect_ratio = tile_img.height / tile_img.width
-    target_tile_height = int(target_tile_width * tile_aspect_ratio)
-    tile_resized = tile_img.resize((target_tile_width, target_tile_height), Image.LANCZOS)
-    tile_np = np.array(tile_resized)
-
-    # Tile texture
-    tile_h, tile_w = tile_np.shape[:2]
-    room_h, room_w = room_np.shape[:2]
-    reps_y = room_h // tile_h + 2
-    reps_x = room_w // tile_w + 2
-    tiled_texture = np.tile(tile_np, (reps_y, reps_x, 1))[:room_h, :room_w]
-    warped_texture = cv2.warpPerspective(tiled_texture, H, (room_w, room_h))
-
-    # Blend
-    room_float = room_np.astype(np.float32) / 255.0
-    texture_float = warped_texture.astype(np.float32) / 255.0
-    room_gray = cv2.cvtColor(room_float, cv2.COLOR_RGB2GRAY)
-    lighting = np.stack([room_gray]*3, axis=-1)
-    lighting = np.clip(lighting * 1.2, 0, 1)
-    lit_texture = np.clip(texture_float * lighting, 0, 1)
-    mask_3ch = np.stack([mask_bin]*3, axis=-1)
-    blended = np.where(mask_3ch == 1, lit_texture, room_float)
-    blended_img = (blended * 255).astype(np.uint8)
-
-    return Image.fromarray(room_np), Image.fromarray(mask_bin * 255), Image.fromarray(warped_texture), Image.fromarray(blended_img)
-
-
-demo = gr.Interface(
-    fn=process,
-    inputs=[gr.Image(label="Room Image", type="pil"), gr.Image(label="Tile Image", type="pil")],
-    outputs=[
-        gr.Image(label="Original Room"),
-        gr.Image(label="Floor Mask"),
-        gr.Image(label="Warped Texture"),
-        gr.Image(label="Final Overlay")
-    ],
-    title="Room Floor Tiler",
-    description="Upload a room image and a tile texture. The floor is automatically detected and overlaid with your selected tile using SegFormer and perspective warping."
-)
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+from torchvision import transforms
+from transformers import SegformerForSemanticSegmentation, AutoImageProcessor
+
+# Device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load model and processor once
+processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b2-finetuned-ade-512-512")
+model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b2-finetuned-ade-512-512").to(device)
+
+def process(room_img, tile_img):
+    room_img = room_img.convert("RGB")
+    tile_img = tile_img.convert("RGB")
+    room_np = np.array(room_img)
+
+    # Segmentation
+    inputs = processor(images=room_img, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    segmentation = outputs.logits.argmax(dim=1).squeeze().cpu().numpy()
+    segmentation_resized = cv2.resize(segmentation.astype(np.uint8), (room_np.shape[1], room_np.shape[0]), interpolation=cv2.INTER_NEAREST)
+
+    # Mask for floor (ADE20K class index 3)
+    floor_class_index = 3
+    mask_bin = (segmentation_resized == floor_class_index).astype(np.uint8)
+
+    # Largest contour
+    contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return room_img, Image.fromarray(mask_bin*255), tile_img, room_img
+    contour = max(contours, key=cv2.contourArea)
+    if len(contour) < 4:
+        return room_img, Image.fromarray(mask_bin*255), tile_img, room_img
+
+    x, y, w, h = cv2.boundingRect(contour)
+    src_pts = np.array([[x, y + h], [x + w, y + h], [x + w, y], [x, y]], dtype=np.float32)
+    offset = h * 0.5
+    dst_pts = np.array([[x - offset, y + h], [x + w + offset, y + h], [x + w, y], [x, y]], dtype=np.float32)
+    H = cv2.getPerspectiveTransform(src_pts, dst_pts)
+
+    # Resize tile
+    target_tile_width = room_np.shape[1] // 10
+    tile_aspect_ratio = tile_img.height / tile_img.width
+    target_tile_height = int(target_tile_width * tile_aspect_ratio)
+    tile_resized = tile_img.resize((target_tile_width, target_tile_height), Image.LANCZOS)
+    tile_np = np.array(tile_resized)
+
+    # Tile texture
+    tile_h, tile_w = tile_np.shape[:2]
+    room_h, room_w = room_np.shape[:2]
+    reps_y = room_h // tile_h + 2
+    reps_x = room_w // tile_w + 2
+    tiled_texture = np.tile(tile_np, (reps_y, reps_x, 1))[:room_h, :room_w]
+    warped_texture = cv2.warpPerspective(tiled_texture, H, (room_w, room_h))
+
+    # Blend
+    room_float = room_np.astype(np.float32) / 255.0
+    texture_float = warped_texture.astype(np.float32) / 255.0
+    room_gray = cv2.cvtColor(room_float, cv2.COLOR_RGB2GRAY)
+    lighting = np.stack([room_gray]*3, axis=-1)
+    lighting = np.clip(lighting * 1.2, 0, 1)
+    lit_texture = np.clip(texture_float * lighting, 0, 1)
+    mask_3ch = np.stack([mask_bin]*3, axis=-1)
+    blended = np.where(mask_3ch == 1, lit_texture, room_float)
+    blended_img = (blended * 255).astype(np.uint8)
+
+    return Image.fromarray(room_np), Image.fromarray(mask_bin * 255), Image.fromarray(warped_texture), Image.fromarray(blended_img)
+
+
+demo = gr.Interface(
+    fn=process,
+    inputs=[gr.Image(label="Room Image", type="pil"), gr.Image(label="Tile Image", type="pil")],
+    outputs=[
+        gr.Image(label="Original Room"),
+        gr.Image(label="Floor Mask"),
+        gr.Image(label="Warped Texture"),
+        gr.Image(label="Final Overlay")
+    ],
+    examples=[
+        ["https://www.thespruce.com/thmb/GtlHim5EsWERYoVi62TnWpu6JTA=/5472x3648/filters:fill(auto,1)/GettyImages-9261821821-5c69c1b7c9e77c0001675a49.jpg", "https://renovlange.de/images/marbles/arabescato.jpg"],
+        ["https://st.hzcdn.com/simgs/57717d160282e776_9-0651/home-design.jpg", "https://renovlange.de/images/marbles/grigio-cenere.jpg"]
+    ],
+    title="Room Floor Tiler",
+    description="Upload a room image and a tile texture. The floor is automatically detected and overlaid with your selected tile using SegFormer and perspective warping."
+)
+
+if __name__ == "__main__":
+    demo.launch()