Update app.py

app.py

@@ -15,34 +15,14 @@ CLIENT = InferenceHTTPClient(
 
 # Set model details
 MODEL_ID = "hvacsym/5"
-IMAGE_PATH = "image1.jpg"
 CONFIDENCE_THRESHOLD = 0.3  # Confidence threshold for filtering predictions
-GRID_SIZE = (3, 3)  # 4x4 segmentation
+GRID_SIZE = (3, 3)  # 3x3 segmentation
 
-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)
-    # Convert back to RGB for colored boxes
-    return image.convert('RGB')
-
-# Ensure image exists before proceeding
-if not os.path.exists(IMAGE_PATH):
-    raise FileNotFoundError(f"Error: The image file '{IMAGE_PATH}' was not found.")
-
-# 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)
+# 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:
@@ -50,114 +30,73 @@ try:
 except:
     font = ImageFont.load_default()
 
-# Dictionary to store total counts
-total_counts = defaultdict(int)
-
-# Colors for boxes
-RED = (255, 0, 0)
-GREEN = (0, 255, 0)
-WHITE = (255, 255, 255)
-BLACK = (0, 0, 0)
-
-# Process each segment
-for row in range(GRID_SIZE[0]):
-    for col in range(GRID_SIZE[1]):
-        # Define segment coordinates
-        x1, y1 = col * seg_w, row * seg_h
-        x2, y2 = (col + 1) * seg_w, (row + 1) * seg_h
-
-        # Crop the segment
-        segment = original_image.crop((x1, y1, x2, y2))
-        draw_segment = ImageDraw.Draw(segment)
-        segment_path = f"/content/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
-        ]
-
-        # Dictionary to count labels in this segment
-        segment_counts = defaultdict(int)
-
-        # Draw bounding boxes on both segment and final image
-        for obj in filtered_predictions:
-            sx, sy, sw, sh = obj["x"], obj["y"], obj["width"], obj["height"]
-            class_name = obj["class"]
-            confidence = obj["confidence"]
-
-            # Update counts
-            segment_counts[class_name] += 1
-            total_counts[class_name] += 1
-
-            # Bounding box coordinates relative to the segment
-            x_min_seg, y_min_seg = sx - sw // 2, sy - sh // 2
-            x_max_seg, y_max_seg = sx + sw // 2, sy + sh // 2
-
-            # Draw on segment with RED
-            draw_segment.rectangle([x_min_seg, y_min_seg, x_max_seg, y_max_seg], outline=RED, width=2)
-
-            # Draw label on segment
-            text = f"{class_name} {confidence:.2f}"
-            text_w, text_h = draw_segment.textbbox((0, 0), text, font=font)[2:]
-            draw_segment.rectangle([x_min_seg, y_min_seg - text_h, x_min_seg + text_w + 4, y_min_seg], fill=BLACK)
-            draw_segment.text((x_min_seg + 2, y_min_seg - text_h), text, fill=WHITE, font=font)
-
-            # Adjust coordinates for the final image
-            x_min_full, y_min_full = x1 + x_min_seg, y1 + y_min_seg
-            x_max_full, y_max_full = x1 + x_max_seg, y1 + y_max_seg
-
-            # Draw on final image with GREEN
-            draw_final.rectangle([x_min_full, y_min_full, x_max_full, y_max_full], outline=GREEN, width=2)
-
-            # Draw label on final image
-            draw_final.rectangle([x_min_full, y_min_full - text_h, x_min_full + text_w + 4, y_min_full], fill=BLACK)
-            draw_final.text((x_min_full + 2, y_min_full - text_h), text, fill=WHITE, font=font)
-
-        # Display the segment with bounding boxes
-        plt.figure(figsize=(5, 5))
-        plt.imshow(segment)  # No need for cmap='gray' as image is now RGB
-        plt.axis("off")
-        plt.title(f"Segment ({row}, {col}) with Detected Symbols")
-        plt.show()
-
-        # Print counts for this segment
-        print(f"Counts in Segment ({row}, {col}):")
-        for label, count in segment_counts.items():
-            print(f"  {label}: {count}")
-        print("-" * 30)
-
-# Display the final image with bounding boxes
-plt.figure(figsize=(10, 10))
-plt.imshow(final_image)  # No need for cmap='gray' as image is now RGB
-plt.axis("off")
-plt.title("Final Image with Detected Symbols")
-plt.show()
-
-# Print total counts for all segments
-print("\nTotal Counts Across All Segments:")
-for label, count in total_counts.items():
-    print(f"{label}: {count}")
+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_final.rectangle([x_min, y_min, x_max, y_max], outline=GREEN, width=2)
+                draw_final.text((x_min, y_min - 10), 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):
-    final_image_path, total_counts = process_image(image_path)  # Calls your existing function
-
-    # Convert count dictionary to readable text
+    """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"),  # Corrected from "file" to "filepath"
+    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."
 )
 
-# Launch the Gradio app
-iface.launch(debug=True)
+iface.launch(debug=True, share=True)