Spaces:
Running on Zero
Running on Zero
| import spaces # must be imported before torch | |
| import torch | |
| import gradio as gr | |
| from PIL import Image, ImageDraw, ImageFont | |
| from transformers import AutoImageProcessor, AutoModelForObjectDetection | |
| MODEL_ID = "joshlyman/rtdetrv2-floorplancad-doors" | |
| processor = AutoImageProcessor.from_pretrained(MODEL_ID) | |
| model = AutoModelForObjectDetection.from_pretrained(MODEL_ID) | |
| model.eval() | |
| COLORS = { | |
| "single_door": "#e6194b", | |
| "double_door": "#3cb44b", | |
| "sliding_door": "#4363d8", | |
| } | |
| def _noop(): # ZeroGPU requires >=1 decorated function; never called (model runs on CPU) | |
| pass | |
| def detect_doors(image: Image.Image, confidence_threshold: float = 0.5): | |
| """Detect single, double, and sliding doors in a floor-plan image and draw the boxes.""" | |
| if image is None: | |
| return None, "Upload a floor-plan image first." | |
| image = image.convert("RGB") | |
| inputs = processor(images=image, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| target_sizes = torch.tensor([image.size[::-1]]) | |
| results = processor.post_process_object_detection( | |
| outputs, threshold=confidence_threshold, target_sizes=target_sizes | |
| )[0] | |
| annotated = image.copy() | |
| draw = ImageDraw.Draw(annotated) | |
| try: | |
| font = ImageFont.truetype("DejaVuSans-Bold.ttf", 16) | |
| except OSError: | |
| font = ImageFont.load_default() | |
| rows = [] | |
| for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]): | |
| label = model.config.id2label[label_id.item()] | |
| conf = score.item() | |
| x0, y0, x1, y1 = [round(v) for v in box.tolist()] | |
| color = COLORS.get(label, "#f58231") | |
| draw.rectangle([x0, y0, x1, y1], outline=color, width=3) | |
| caption = f"{label} {conf:.2f}" | |
| text_bbox = draw.textbbox((x0, y0), caption, font=font) | |
| draw.rectangle( | |
| [x0, text_bbox[1] - 2, text_bbox[2] + 4, text_bbox[3] + 2], fill=color | |
| ) | |
| draw.text((x0 + 2, y0), caption, fill="white", font=font) | |
| rows.append([label, f"{conf:.2%}", f"({x0}, {y0}, {x1}, {y1})"]) | |
| if not rows: | |
| summary = "No doors detected above this confidence threshold." | |
| else: | |
| summary = f"{len(rows)} door(s) detected." | |
| return annotated, gr.Dataframe( | |
| headers=["type", "confidence", "box (x0, y0, x1, y1)"], value=rows | |
| ), summary | |
| with gr.Blocks(title="Floor Plan Door Detector") as demo: | |
| gr.Markdown( | |
| "# Floor Plan Door Detector\n" | |
| f"RT-DETRv2 fine-tuned on floor-plan CAD drawings — " | |
| f"[{MODEL_ID}](https://huggingface.co/{MODEL_ID}) — " | |
| "detects **single**, **double**, and **sliding** doors. Upload a floor-plan " | |
| "image (or an exported CAD/PDF page as an image) to see the detections." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_input = gr.Image(type="pil", label="Floor plan") | |
| threshold = gr.Slider( | |
| minimum=0.05, maximum=0.95, value=0.5, step=0.05, | |
| label="Confidence threshold", | |
| ) | |
| run_btn = gr.Button("Detect doors", variant="primary") | |
| with gr.Column(): | |
| image_output = gr.Image(type="pil", label="Detections") | |
| summary_output = gr.Textbox(label="Summary", interactive=False) | |
| table_output = gr.Dataframe( | |
| headers=["type", "confidence", "box (x0, y0, x1, y1)"], | |
| label="Detected doors", | |
| ) | |
| run_btn.click( | |
| fn=detect_doors, | |
| inputs=[image_input, threshold], | |
| outputs=[image_output, table_output, summary_output], | |
| api_name="detect_doors", | |
| ) | |
| image_input.change( | |
| fn=detect_doors, | |
| inputs=[image_input, threshold], | |
| outputs=[image_output, table_output, summary_output], | |
| ) | |
| demo.launch(mcp_server=True) | |