Spaces:
Running
Running
File size: 4,057 Bytes
09ebd0a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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()
|