Update app.py

app.py

@@ -3,10 +3,10 @@ import gradio as gr
 from PIL import Image
 import numpy as np
 
-# Load the YOLOv5 model
+# Load the YOLOv5 model (adjust the path to your model if needed)
 model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5s.pt')  # Adjust if needed
 
-# Example function to calculate materials based on detected areas (you can customize the formulas as needed)
+# Example function to calculate materials based on detected areas
 def calculate_materials(detected_objects, image_width, image_height):
     materials = {
         "cement": 0,
@@ -40,8 +40,37 @@ def calculate_materials(detected_objects, image_width, image_height):
 
 # Define the function for image inference
 def predict_image(image):
-    results = model(image)  # Run inference on the input image
-    detected_objects = results.pandas().xywh[0]  # Get the detected objects as pandas dataframe
+    # Run inference on the input image
+    results = model(image)
     
-    # Calculate real
+    # Get the detected objects as pandas dataframe (xywh format)
+    detected_objects = results.pandas().xywh[0]  # First image in batch
+    
+    # Print out the detection results for debugging purposes
+    print(f"Detected objects: {detected_objects}")
+    
+    # Assume blueprint image size (in cm, adjust based on real-world image size)
+    image_width = 91  # Example width in cm (adjust this)
+    image_height = 61  # Example height in cm (adjust this)
+    
+    # Process the detected objects and calculate materials
+    detected_objects_list = []
+    for _, row in detected_objects.iterrows():
+        detected_objects_list.append({
+            'name': row['name'],   # Detected object class name (e.g., 'wall', 'foundation')
+            'bbox': [row['xmin'], row['ymin'], row['xmax'], row['ymax']]  # Bounding box coordinates
+        })
+    
+    # Calculate materials based on detected objects
+    materials = calculate_materials(detected_objects_list, image_width, image_height)
+    
+    # Return the materials as a dictionary
+    return materials
+
+# Set up Gradio interface for image input and JSON output
+interface = gr.Interface(fn=predict_image, inputs=gr.Image(), outputs=gr.JSON())
+
+# Launch the Gradio interface
+if __name__ == "__main__":
+    interface.launch()