Update app.py

app.py

@@ -13,35 +13,17 @@ model.eval()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-# Define categories for construction activities and materials
-construction_terms = {
-    "activities": ["pouring", "scaffolding", "building", "excavation", "piling", "digging", "cementing", "welding", "cutting", "assembling", "drilling"],
-    "materials": ["concrete", "steel", "wood", "bricks", "cement", "sand", "mortar", "rebar", "plaster", "tiles"],
-    "progress": ["completed", "ongoing", "in-progress", "starting", "finished", "under construction"]
-}
-
-# Function to detect activities and materials
-def detect_construction_info(caption):
-    activity_found = []
-    material_found = []
-    progress_found = []
-
-    # Split the caption into words and check for the terms
-    for word in caption.split():
-        word_lower = word.lower()
-        if word_lower in construction_terms["activities"]:
-            activity_found.append(word)
-        elif word_lower in construction_terms["materials"]:
-            material_found.append(word)
-        elif word_lower in construction_terms["progress"]:
-            progress_found.append(word)
+# Inference function to generate captions from images dynamically
+def generate_captions_from_image(image):
+    if image.mode != "RGB":
+        image = image.convert("RGB")
     
-    # Build the informative output
-    activity_str = ", ".join(activity_found) if activity_found else "No specific activities detected."
-    material_str = ", ".join(material_found) if material_found else "No materials detected."
-    progress_str = ", ".join(progress_found) if progress_found else "No progress information available."
-
-    return f"Activities: {activity_str}\nMaterials: {material_str}\nProgress: {progress_str}"
+    # 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
 def generate_dpr(files):
@@ -59,16 +41,11 @@ def generate_dpr(files):
         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)
-
-        # Get detailed construction information based on the caption
-        detailed_caption = detect_construction_info(caption)
+        # Dynamically generate a caption based on the image
+        caption = generate_captions_from_image(image)
         
-        # Generate DPR section for this image
-        dpr_section = f"\nImage: {file.name}\n{detailed_caption}\n"
+        # Generate DPR section for this image with dynamic caption
+        dpr_section = f"\nImage: {file.name}\nDescription: {caption}\n"
         dpr_text.append(dpr_section)
 
     # Generate a PDF report
@@ -77,7 +54,7 @@ def generate_dpr(files):
     c.drawString(100, 750, "Daily Progress Report")
     c.drawString(100, 730, f"Generated on: {current_time}")
 
-    # Add the detailed captions for each image to the PDF
+    # Add the detailed captions for each image to the PDF (in text format)
     y_position = 700
     for section in dpr_text:
         c.drawString(100, y_position, section)
@@ -96,7 +73,7 @@ iface = gr.Interface(
     inputs=gr.Files(type="filepath", label="Upload Site Photos"),  # Handle batch upload of images
     outputs="file",
     title="Daily Progress Report Generator",
-    description="Upload up to 10 site photos. The AI model will detect construction activities, materials, and progress and generate a PDF report.",
+    description="Upload up to 10 site photos. The AI model will dynamically detect construction activities, materials, and progress and generate a PDF report.",
     allow_flagging="never"  # Optional: Disable flagging
 )