File size: 4,424 Bytes
0facd8f f3e768f 0facd8f f3e768f 0facd8f c9068cc 0facd8f f3e768f c9068cc f3e768f c9068cc f3e768f 0facd8f f3e768f 0facd8f f3e768f 0facd8f f3e768f |
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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
from inference_sdk import InferenceHTTPClient
from PIL import Image, ImageDraw, ImageFont, ImageEnhance
import matplotlib.pyplot as plt
import os
import gradio as gr
from collections import defaultdict
API_KEY = os.getenv("ROBOFLOW_API_KEY")
# Initialize the Roboflow client
CLIENT = InferenceHTTPClient(
api_url="https://detect.roboflow.com",
api_key=API_KEY
)
# Set model details
MODEL_ID = "hvacsym/5"
CONFIDENCE_THRESHOLD = 0.3 # Confidence threshold for filtering predictions
GRID_SIZE = (3, 3) # 3x3 segmentation
# Colors for bounding boxes
RED = (255, 0, 0)
GREEN = (0, 255, 0)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
# Load font for labeling
try:
font = ImageFont.truetype("arial.ttf", 14)
except:
font = ImageFont.load_default()
def enhance_image(image):
"""Enhance image by adjusting brightness and contrast."""
if image.mode != 'L':
image = image.convert('L')
brightness = ImageEnhance.Brightness(image)
image = brightness.enhance(1.3)
contrast = ImageEnhance.Contrast(image)
image = contrast.enhance(1.2)
return image.convert('RGB') # Convert back to RGB for colored boxes
def process_image(image_path):
"""Processes an image by running inference and drawing bounding boxes."""
# Load and enhance the original image
original_image = Image.open(image_path)
original_image = enhance_image(original_image)
width, height = original_image.size
seg_w, seg_h = width // GRID_SIZE[1], height // GRID_SIZE[0]
# Create a copy of the full image to draw bounding boxes
final_image = original_image.copy()
draw_final = ImageDraw.Draw(final_image)
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 on the segment
result = CLIENT.infer(segment_path, model_id=MODEL_ID)
# Filter predictions based on confidence
filtered_predictions = [
pred for pred in result["predictions"] if pred["confidence"] * 100 >= CONFIDENCE_THRESHOLD
]
# Draw bounding boxes and count labels
for obj in filtered_predictions:
class_name = obj["class"]
total_counts[class_name] += 1
x_min, y_min = x1 + obj["x"] - obj["width"] // 2, y1 + obj["y"] - obj["height"] // 2
x_max, y_max = x1 + obj["x"] + obj["width"] // 2, y1 + obj["y"] + obj["height"] // 2
# Draw bounding box
draw_final.rectangle([x_min, y_min, x_max, y_max], outline=GREEN, width=2)
# Draw extended label above the bounding box
text_size = draw_final.textbbox((0, 0), class_name, font=font)
text_width = text_size[2] - text_size[0]
text_height = text_size[3] - text_size[1]
text_x = x_min
text_y = y_min - text_height - 5 if y_min - text_height - 5 > 0 else y_max + 5
draw_final.rectangle([text_x, text_y, text_x + text_width + 6, text_y + text_height + 2], fill=BLACK)
draw_final.text((text_x + 3, text_y), class_name, fill=WHITE, font=font)
# Save the final processed image
final_image_path = "processed_image.png"
final_image.save(final_image_path)
return final_image_path, total_counts
def process_uploaded_image(image_path):
"""Handles uploaded image and processes it."""
final_image_path, total_counts = process_image(image_path)
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"),
outputs=[gr.Image(type="filepath"), gr.Text()],
title="HVAC Symbol Detector",
description="Upload an HVAC blueprint image. The model will segment it, detect symbols, and return the final image with bounding boxes along with symbol counts."
)
iface.launch(debug=True, share=True) |