Update app.py

app.py

@@ -1,29 +1,13 @@
-from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
 import gradio as gr
 import torch
 from datetime import datetime
+from ultralytics import YOLO
 
-# Load BLIP model and processor
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-model.eval()
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# Load YOLOv8 model (trained on construction dataset)
+model = YOLO('yolov8n.pt')  # Path to pre-trained model on construction dataset
 
-# Inference function to generate captions dynamically based on image content
-def generate_captions_from_image(image):
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    
-    # Preprocess the image and generate a caption
-    inputs = processor(image, return_tensors="pt").to(device, torch.float16)
-    output = model.generate(**inputs, max_new_tokens=50)
-    caption = processor.decode(output[0], skip_special_tokens=True)
-    
-    return caption
-
-# Function to generate the daily progress report (DPR) text
+# Function to generate DPR text based on detections
 def generate_dpr(files):
     dpr_text = []
     current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
@@ -34,18 +18,44 @@ def generate_dpr(files):
     # Process each uploaded file (image)
     for file in files:
         # Open the image from file path
-        image = Image.open(file.name)  # Using file.name for filepath
+        image = Image.open(file.name)
+        
+        # Perform object detection with YOLOv8
+        results = model(image)  # Perform detection
+        
+        # Parse detections (activities, materials, etc.)
+        detected_objects = results.names  # Object names detected by the model
+        detections = results.pandas().xywh  # Get the dataframe with detection results
+        
+        detected_activities = []
+        detected_materials = []
         
-        if image.mode != "RGB":
-            image = image.convert("RGB")
+        # Define construction activity and material categories
+        construction_activities = ['scaffolding', 'concrete pouring', 'welding', 'excavation']
+        construction_materials = ['concrete', 'steel', 'bricks', 'cement', 'sand']
         
-        # Dynamically generate a caption based on the image
-        caption = generate_captions_from_image(image)
+        # Check the detected objects and categorize them
+        for obj in detected_objects:
+            if obj.lower() in construction_activities:
+                detected_activities.append(obj)
+            elif obj.lower() in construction_materials:
+                detected_materials.append(obj)
+        
+        # Build a detailed report for this image
+        dpr_section = f"\nImage: {file.name}\n"
+        
+        if detected_activities:
+            dpr_section += f"Detected Activities: {', '.join(detected_activities)}\n"
+        else:
+            dpr_section += "No construction activities detected.\n"
+        
+        if detected_materials:
+            dpr_section += f"Detected Materials: {', '.join(detected_materials)}\n"
+        else:
+            dpr_section += "No materials detected.\n"
         
-        # Generate DPR section for this image with dynamic caption
-        dpr_section = f"\nImage: {file.name}\nDescription: {caption}\n"
         dpr_text.append(dpr_section)
-
+    
     # Return the generated DPR as a text output
     return "\n".join(dpr_text)
 
@@ -55,9 +65,8 @@ iface = gr.Interface(
     inputs=gr.Files(type="filepath", label="Upload Site Photos"),  # Handle batch upload of images
     outputs="text",  # Display the DPR as text in the output section
     title="Daily Progress Report Generator",
-    description="Upload up to 10 site photos. The AI model will dynamically detect construction activities, materials, and progress and generate a text-based Daily Progress Report (DPR).",
+    description="Upload up to 10 site photos. The AI model will detect construction activities, materials, and progress and generate a text-based Daily Progress Report (DPR).",
     allow_flagging="never"  # Optional: Disable flagging
 )
 
 iface.launch()
-