Update app.py

app.py

@@ -1,12 +1,14 @@
 import torch
-import gradio as gr
+from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
 from PIL import Image
+import gradio as gr
 import numpy as np
 
-# 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
+# Load the pre-trained DETR model
+model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+extractor = AutoFeatureExtractor.from_pretrained("facebook/detr-resnet-50")
 
-# Define the function to calculate materials based on detected areas
+# Function to calculate materials based on detected areas (example: walls and foundations)
 def calculate_materials(detected_objects, image_width, image_height):
     materials = {
         "cement": 0,
@@ -14,63 +16,62 @@ def calculate_materials(detected_objects, image_width, image_height):
         "steel": 0
     }
 
-    # Proportionality factors (simplified for this example, adjust based on real-world data)
+    # Proportionality factors (simplified, adjust based on real-world data)
     for obj in detected_objects:
-        # Get the bounding box coordinates
+        # Bounding box coordinates in the format [xmin, ymin, xmax, ymax]
         x1, y1, x2, y2 = obj['bbox']
-
-        # Convert the bounding box coordinates to real-world units based on image size and blueprint size
-        width = (x2 - x1) * image_width  # Convert to real-world width (in cm or meters)
-        height = (y2 - y1) * image_height  # Convert to real-world height (in cm or meters)
         
-        # Calculate the area (length × width) in real-world units
-        area = width * height  # cm² or m² based on the scale
-
-        # Debugging output to verify bounding box size
-        print(f"Detected {obj['name']} with area {area} cm²")  # Adjust units based on the scale
+        # Calculate real-world dimensions (assuming you have a known scale for your blueprint)
+        width = (x2 - x1) * image_width  # Convert to real-world width
+        height = (y2 - y1) * image_height  # Convert to real-world height
+        
+        # Calculate area
+        area = width * height  # cm² or m² depending on your scale
 
-        # Example material estimation based on detected object type
+        # Print area for debugging
+        print(f"Detected {obj['name']} with area {area} cm²")
+        
+        # Material estimation based on the object name
         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
         
         elif 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
+            materials['cement'] += area * 0.2
+            materials['bricks'] += area * 15
+            materials['steel'] += area * 0.1
 
     return materials
 
 # Define the function for image inference
 def predict_image(image):
-    # Run inference on the input image
-    results = model(image)
-    
-    # Get the detected objects as pandas dataframe (xywh format)
-    detected_objects = results.pandas().xywh[0]  # First image in batch
-    
-    # Debugging output: Print out the detected objects for inspection
-    print(f"Detected objects: {detected_objects}")
-    
-    # Set the confidence threshold (e.g., 0.5 means 50% confidence)
-    confidence_threshold = 0.5
-    detected_objects = detected_objects[detected_objects['confidence'] > confidence_threshold]
+    # Convert image to the required format for the model
+    inputs = extractor(images=image, return_tensors="pt")
+
+    # Run inference with DETR
+    outputs = model(**inputs)
     
-    # Assume blueprint image size (in cm, adjust based on real-world image size)
-    image_width = 91  # Example width in cm (adjust this to the real-world blueprint size)
-    image_height = 61  # Example height in cm (adjust this to the real-world blueprint size)
+    # Get the predictions from the output (boxes and labels)
+    target_sizes = torch.tensor([image.size[::-1]])  # height, width
+    results = extractor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.5)[0]
+
+    detected_objects = []
     
-    # 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
+    # Process the detected objects and extract bounding boxes and class names
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = box.tolist()  # Convert box to list
+        detected_objects.append({
+            'name': label.item(),  # Get the class name
+            'bbox': box  # Bounding box [xmin, ymin, xmax, ymax]
         })
+
+    # Assume blueprint image size (adjust this based on your image scale)
+    image_width = image.size[0]  # Image width in pixels
+    image_height = image.size[1]  # Image height in pixels
     
-    # Calculate materials based on detected objects
-    materials = calculate_materials(detected_objects_list, image_width, image_height)
+    # Calculate materials based on the detected objects
+    materials = calculate_materials(detected_objects, image_width, image_height)
     
     # Return the materials as a dictionary
     return materials