Update app.py

app.py

@@ -1,15 +1,15 @@
 import torch
-from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
-from PIL import Image
 import gradio as gr
+from PIL import Image
+from pdf2image import convert_from_path
 import numpy as np
+from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
 
 # Load the pre-trained DETR model and feature extractor
 model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 extractor = AutoFeatureExtractor.from_pretrained("facebook/detr-resnet-50")
 
 # Simplified COCO label mapping (DETR uses COCO dataset)
-# In a real scenario, you'd need a custom model trained for construction labels
 COCO_LABELS = {56: "wall", 60: "foundation"}  # Pretending chair (56) is "wall", dining table (60) is "foundation"
 
 # Function to calculate materials based on detected areas
@@ -39,13 +39,20 @@ def calculate_materials(detected_objects, image_width, image_height):
 
     return materials
 
-# Define the function for image inference
-def predict_image(image):
-    # Convert Gradio image (numpy array) to PIL Image
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
+# Function to process PDFs and convert to images
+def pdf_to_image(pdf_file):
+    images = convert_from_path(pdf_file, first_page=1, last_page=1)  # Convert the first page of the PDF
+    return images[0]  # Return the first page as an image
 
-    # Prepare image for the model
+# Define the function for image inference
+def predict_image(file):
+    # Check if the input file is a PDF or image
+    if isinstance(file, str) and file.endswith('.pdf'):
+        image = pdf_to_image(file)  # Convert PDF to image
+    else:
+        image = Image.open(file)  # Open the image file
+
+    # Convert the image to the correct format for the model
     inputs = extractor(images=image, return_tensors="pt")
 
     # Run inference with DETR
@@ -59,7 +66,6 @@ def predict_image(image):
     for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
         box = box.tolist()
         label_id = label.item()
-        # Map the label to a construction-related name (simplified for demo)
         label_name = COCO_LABELS.get(label_id, "unknown")
         if label_name != "unknown":  # Only process relevant objects
             detected_objects.append({
@@ -83,12 +89,12 @@ def predict_image(image):
 # Set up Gradio interface
 interface = gr.Interface(
     fn=predict_image,
-    inputs=gr.Image(type="numpy", label="Upload Blueprint Image"),
+    inputs=gr.File(label="Upload Blueprint (PDF or Image)"),
     outputs=gr.JSON(label="Material Estimates"),
     title="Blueprint Material Estimator",
-    description="Upload a blueprint image to estimate construction materials."
+    description="Upload a blueprint image or PDF to estimate construction materials."
 )
 
 # Launch the interface
 if __name__ == "__main__":
-    interface.launch(share=False)
+    interface.launch(share=False)