Spaces:
Sleeping
Sleeping
File size: 3,474 Bytes
ba7630e cb97926 ba7630e |
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 95 96 97 98 99 |
import gradio as gr
from inference_sdk import InferenceHTTPClient
from PIL import Image, ImageDraw, ImageFont
import os
from collections import defaultdict
API_KEY = os.getenv("ROBOFLOW_API_KEY")
CLIENT = InferenceHTTPClient(
api_url="https://detect.roboflow.com",
api_key=API_KEY
)
# Model settings
MODEL_ID = "hvacsym/5"
CONFIDENCE_THRESHOLD = 0.3 # Confidence threshold for filtering predictions
GRID_SIZE = (3, 3) # 4x4 segmentation
def process_image(image_path):
"""Processes an uploaded image and returns the final image with bounding boxes & symbol counts."""
original_image = Image.open(image_path)
width, height = original_image.size
seg_w, seg_h = width // GRID_SIZE[1], height // GRID_SIZE[0]
# Create a copy of the image for bounding boxes
final_image = original_image.copy()
draw_final = ImageDraw.Draw(final_image)
# Load font
try:
font = ImageFont.truetype("arial.ttf", 10)
except:
font = ImageFont.load_default()
# Dictionary for total counts
total_counts = defaultdict(int)
# Process each segment
for row in range(GRID_SIZE[0]):
for col in range(GRID_SIZE[1]):
x1, y1 = col * seg_w, row * seg_h
x2, y2 = (col + 1) * seg_w, (row + 1) * seg_h
segment = original_image.crop((x1, y1, x2, y2))
segment_path = f"segment_{row}_{col}.png"
segment.save(segment_path)
# Run inference
result = CLIENT.infer(segment_path, model_id=MODEL_ID)
filtered_predictions = [
pred for pred in result["predictions"] if pred["confidence"] * 100 >= CONFIDENCE_THRESHOLD
]
# Draw bounding boxes and update counts
for obj in filtered_predictions:
sx, sy, sw, sh = obj["x"], obj["y"], obj["width"], obj["height"]
class_name = obj["class"]
confidence = obj["confidence"]
total_counts[class_name] += 1
# Adjust coordinates for final image
x_min, y_min = x1 + (sx - sw // 2), y1 + (sy - sh // 2)
x_max, y_max = x1 + (sx + sw // 2), y1 + (sy + sh // 2)
# Draw bounding box
draw_final.rectangle([x_min, y_min, x_max, y_max], outline="green", width=2)
# Draw label
text = f"{class_name} {confidence:.2f}"
draw_final.text((x_min + 2, y_min - 10), text, fill="white", font=font)
# Save final image with bounding boxes
final_image_path = "final_detected_image.png"
final_image.save(final_image_path)
return final_image_path, total_counts
def process_uploaded_image(image):
"""Gradio wrapper function that calls `process_image` and formats the output."""
final_image_path, total_counts = process_image(image)
# Convert count dictionary to readable text
count_text = "\n".join([f"{label}: {count}" for label, count in total_counts.items()])
return final_image_path, count_text
# Deploy with Gradio
iface = gr.Interface(
fn=process_uploaded_image,
inputs=gr.Image(type="filepath"), # Gradio input expects a file path
outputs=[gr.Image(type="filepath"), gr.Text()],
title="HVAC Symbol Detector",
description="Upload an HVAC blueprint image. The model will detect symbols and return the final image with bounding boxes along with symbol counts."
)
# Launch the Gradio app
iface.launch()
|