Kushalmanda's picture
Update app.py
05f6e4f verified
raw
history blame
1.98 kB
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