Spaces:
Build error
Build error
import torch | |
import gradio as gr | |
from PIL import Image | |
import numpy as np | |
# Load the YOLOv5 model | |
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) | |
def calculate_materials(detected_objects, image_width, image_height): | |
materials = { | |
"cement": 0, | |
"bricks": 0, | |
"steel": 0 | |
} | |
# Proportionality factors (simplified for this example, adjust based on real-world data) | |
for obj in detected_objects: | |
# Calculate bounding box area in real-world units (cm or meters, as per the blueprint size) | |
x1, y1, x2, y2 = obj['bbox'] # Coordinates of the bounding box | |
width = (x2 - x1) * image_width # Convert to real-world width | |
height = (y2 - y1) * image_height # Convert to real-world height | |
# Calculate the area (length × width) | |
area = width * height # Simplified area calculation | |
print(f"Detected {obj['name']} with area {area} cm²") # Debugging output | |
if obj['name'] == 'wall': # Example: For 'wall' objects | |
materials['cement'] += area * 0.1 # Cement estimation (in kg) | |
materials['bricks'] += area * 10 # Bricks estimation | |
materials['steel'] += area * 0.05 # Steel estimation | |
if obj['name'] == 'foundation': # Example: For 'foundation' objects | |
materials['cement'] += area * 0.2 # More cement for foundation | |
materials['bricks'] += area * 15 # More bricks for foundation | |
materials['steel'] += area * 0.1 # More steel for foundation | |
return materials | |
# 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 | |
# Calculate real | |