Update app.py

app.py

@@ -1,18 +1,54 @@
 import torch
 import gradio as gr
 from PIL import Image
+import numpy as np
 
 # Load the YOLOv5 model from the uploaded file (e.g., 'yolov5s.pt')
 model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt')  # Adjust the file name if needed
 
-# Define the function for image inference
+# Define the function for image inference and cement calculation
 def predict_image(image):
-    results = model(image)  # Run inference on the input image
-    results.show()  # Optionally visualize the results (bounding boxes)
-    return results.pandas().xywh  # Return the results (bounding box coordinates)
+    # Run inference on the input image
+    results = model(image)
+    
+    # Get the bounding boxes and class names
+    predictions = results.pandas().xywh
+    
+    # Filter for bricks (assuming the model detects a "brick" class)
+    brick_class = 'brick'  # Define the brick class label used in the YOLO model
+    brick_data = predictions[predictions['name'] == brick_class]
+    
+    # If bricks are detected, calculate cement estimation
+    if len(brick_data) > 0:
+        # Predefined brick dimensions in cm (assuming a standard brick size of 20x10x5 cm)
+        brick_length = 20  # cm
+        brick_width = 10   # cm
+        brick_height = 5   # cm
+        
+        # Assuming a cement-to-brick ratio (this depends on the specific mortar mixture used)
+        cement_per_brick = 0.01  # (in cubic meters, adjust this based on your own research or formula)
+        
+        # Total bricks detected
+        total_bricks = len(brick_data)
+        
+        # Calculate the volume of one brick in cubic meters (converted from cm)
+        brick_volume = (brick_length / 100) * (brick_width / 100) * (brick_height / 100)
+        
+        # Calculate the total volume of all bricks in cubic meters
+        total_brick_volume = brick_volume * total_bricks
+        
+        # Estimate the cement required in cubic meters based on the volume
+        total_cement = cement_per_brick * total_bricks
+        
+        # Output the results
+        results_text = f"Detected {total_bricks} bricks.\nTotal Brick Volume: {total_brick_volume:.2f} m³\nEstimated Cement Needed: {total_cement:.2f} m³"
+        
+        return results_text, results.show()  # Optionally visualize the results (bounding boxes)
+    else:
+        return "No bricks detected in the image.", results.show()
 
 # Set up Gradio interface to allow image uploads and get predictions
-interface = gr.Interface(fn=predict_image, inputs=gr.Image(), outputs=gr.Dataframe())
+interface = gr.Interface(fn=predict_image, inputs=gr.Image(), outputs=["text", "image"])
 
 # Launch the interface
 interface.launch()