Update app.py

app.py

@@ -1,108 +1,117 @@
 from flask import Flask, request, render_template, send_file
 from inference_sdk import InferenceHTTPClient
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, ImageDraw, ImageFont, ImageEnhance
 import os
 from collections import defaultdict
 
-# Define a writable directory for saving image segments
-SEGMENT_DIR = "static/segments"
-
-# Ensure the "static/segments" directory exists
-if not os.path.exists(SEGMENT_DIR):
-    os.makedirs(SEGMENT_DIR)
-
 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
+    api_key=API_KEY  # Secure way to use API key
 )
 
 # Model settings
 MODEL_ID = "hvacsym/5"
-CONFIDENCE_THRESHOLD = 30  # Confidence threshold for filtering predictions
-GRID_SIZE = (4, 4)  # 4x4 segmentation
-
+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 uploaded image and returns the final image with bounding boxes & symbol counts."""
+    """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 image for bounding boxes
+    
+    # Create a copy of the full image to draw bounding boxes
     final_image = original_image.copy()
     draw_final = ImageDraw.Draw(final_image)
-
-    # Load font
-    try:
-        font = ImageFont.truetype("arial.ttf", 15)
-    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 = os.path.join(SEGMENT_DIR, f"segment_{row}_{col}.png")
+            segment_path = f"static/processed/segment_{row}_{col}.png"
+            segment.save(segment_path)
             
-            segment.save(segment_path)  # ✅ Now saving in a writable folder
-
-            # Run inference
+            # 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 update counts
+            
+            # Draw bounding boxes and count labels
             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)
-
+                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 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 = "static/final_detected_image.png"
+                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 the static directory exists
-            static_dir = "static"
-            os.makedirs(static_dir, exist_ok=True)  # ✅ Create if it doesn't exist
+            # Ensure static/processed directory exists
+            os.makedirs(PROCESSED_DIR, exist_ok=True)
 
-            image_path = os.path.join(static_dir, "uploaded_image.png")
-
-            image_file.save(image_path)  # ✅ Save file without modifying permissions
+            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)
@@ -111,6 +120,5 @@ def index():
 
     return render_template("index.html", final_image=None, counts=None)
 
-
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)