Spaces:
Sleeping
Sleeping
File size: 4,783 Bytes
47526b2 25030bd 47526b2 25030bd 47526b2 25030bd 47526b2 25030bd c4f3e32 47526b2 25030bd 47526b2 25030bd 47526b2 25030bd 47526b2 25030bd 47526b2 25030bd c4f3e32 25030bd 47526b2 25030bd 47526b2 25030bd 47526b2 25030bd 47526b2 25030bd 47526b2 25030bd 0396dd2 25030bd 47526b2 |
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 114 115 116 117 118 119 120 121 122 123 124 125 |
from flask import Flask, request, render_template, send_file
from inference_sdk import InferenceHTTPClient
from PIL import Image, ImageDraw, ImageFont, ImageEnhance
import os
from collections import defaultdict
app = Flask(__name__)
# Ensure the "static/processed" directory exists for output images
PROCESSED_DIR = "static/processed"
if not os.path.exists(PROCESSED_DIR):
os.makedirs(PROCESSED_DIR, exist_ok=True)
# Securely get API key from Hugging Face Secrets
API_KEY = os.getenv("ROBOFLOW_API_KEY")
# Initialize the Roboflow client
CLIENT = InferenceHTTPClient(
api_url="https://detect.roboflow.com",
api_key=API_KEY # Secure way to use API key
)
# Model settings
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"static/processed/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 = "static/processed/processed_image.png"
final_image.save(final_image_path)
return final_image_path, total_counts
@app.route("/", methods=["GET", "POST"])
def index():
if request.method == "POST":
image_file = request.files["image"]
if image_file:
# Ensure static/processed directory exists
os.makedirs(PROCESSED_DIR, exist_ok=True)
image_path = os.path.join(PROCESSED_DIR, "uploaded_image.png")
image_file.save(image_path)
# Process image
final_image_path, total_counts = process_image(image_path)
return render_template("index.html", final_image=final_image_path, counts=total_counts)
return render_template("index.html", final_image=None, counts=None)
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)
|