Update app.py

app.py

@@ -4,15 +4,25 @@ import numpy as np
 from PIL import Image
 import torch
 from transformers import DetrImageProcessor, DetrForObjectDetection
+import os
 
 # Load DETR model for object detection
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+# We use try-except to handle potential model loading issues in constrained environments
+try:
+    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
+    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    processor = None
+    model = None
 
 def detect_objects(image):
     """
     Detect objects in an image using DETR model
     """
+    if model is None or processor is None:
+        return image, "Model failed to load. Please check logs."
+
     # Convert PIL image to numpy array
     if isinstance(image, Image.Image):
         image = np.array(image)
@@ -25,7 +35,9 @@ def detect_objects(image):
     
     # Prepare image for the model
     inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
     
     # Post-process outputs
     target_sizes = torch.tensor([image.shape[:2]])
@@ -83,6 +95,9 @@ def process_image(image, mode):
     """
     Process image based on selected mode
     """
+    if image is None:
+        return None, "Please upload an image."
+        
     if mode == "Object Detection":
         return detect_objects(image)
     else:  # Edge Detection
@@ -90,7 +105,7 @@ def process_image(image, mode):
         return edges, "Edge detection completed"
 
 # Create Gradio interface
-with gr.Blocks(title="Object Scanner") as demo:
+with gr.Blocks(title="Object Scanner", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🔍 Object Scanner with OpenCV & Hugging Face")
     gr.Markdown("Upload an image to detect objects or scan edges using computer vision")
     
@@ -109,16 +124,17 @@ with gr.Blocks(title="Object Scanner") as demo:
             output_text = gr.Markdown(label="Detection Results")
     
     # Examples
-    gr.Examples(
-        examples=[
-            ["examples/sample1.jpg", "Object Detection"],
-            ["examples/sample2.jpg", "Edge Detection"],
-        ],
-        inputs=[input_image, mode],
-        outputs=[output_image, output_text],
-        fn=process_image,
-        cache_examples=False,
-    )
+    if os.path.exists("examples"):
+        gr.Examples(
+            examples=[
+                ["examples/sample1.jpg", "Object Detection"],
+                ["examples/sample2.jpg", "Edge Detection"],
+            ],
+            inputs=[input_image, mode],
+            outputs=[output_image, output_text],
+            fn=process_image,
+            cache_examples=False,
+        )
     
     scan_btn.click(
         fn=process_image,
@@ -127,4 +143,4 @@ with gr.Blocks(title="Object Scanner") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()