Update app.py

app.py

@@ -1,11 +1,10 @@
 import os
 import gradio as gr
-from langchain_google_genai.chat_models import ChatGoogleGenerativeAI  # Import for Gemini
+from transformers import DetrForObjectDetection, DetrImageProcessor
 from PIL import Image
-import json
-from transformers import DetrImageProcessor, DetrForObjectDetection
 import torch
 import requests
+import json
 
 # Load credentials (stringified JSON) from environment variable
 credentials_string = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
@@ -25,10 +24,9 @@ os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "service_account.json"
 # Initialize Gemini model
 llm = ChatGoogleGenerativeAI(model='gemini-1.5-pro')
 
-# Initialize DETR model and processor
+# Load the model and processor for DETR
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
-model.eval()
 
 # Global chat history variable
 chat_history = []
@@ -42,37 +40,28 @@ def chat_with_gemini(message):
 
 def analyze_image(image_path):
     global chat_history
-
-    # Load and preprocess image
-    image = Image.open(image_path).convert("RGB")
-    inputs = processor(images=image, return_tensors="pt")
-
-    # Inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    # Extract predictions
-    logits = outputs.logits
-    boxes = outputs.pred_boxes
-
-    # Filter predictions by high confidence scores
-    scores = logits.softmax(-1)[0, :, :-1].max(-1).values
-    high_scores_indices = scores > 0.9  # Adjust the threshold as needed
-    predicted_classes = logits.softmax(-1)[0, high_scores_indices, :-1].argmax(-1)
-    predicted_boxes = boxes[0, high_scores_indices].tolist()
-
-    # Map class IDs to labels
-    labels = [processor.config.id2label[idx.item()] for idx in predicted_classes]
-
-    # Combine predictions
-    predictions = [{"label": label, "box": box} for label, box in zip(labels, predicted_boxes)]
-
-    # Create response
-    if predictions:
-        detected_objects = ', '.join([p["label"] for p in predictions])
-        bot_response = f"The image contains: {detected_objects}."
-    else:
-        bot_response = "No objects with high confidence were detected."
+    try:
+        # Open the image file
+        image = Image.open(image_path).convert("RGB")
+        # Preprocess the image for DETR
+        inputs = processor(images=image, return_tensors="pt")
+
+        # Perform inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        # Get predictions
+        target_sizes = torch.tensor([image.size[::-1]])  # (height, width)
+        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+
+        # If any objects detected, display labels
+        if len(results["labels"]) > 0:
+            bot_response = f"Objects detected: {', '.join(map(str, results['labels'].tolist()))}."
+        else:
+            bot_response = "No objects detected."
+
+    except Exception as e:
+        bot_response = f"Error processing the image: {str(e)}"
 
     chat_history.append(("Uploaded an image for analysis", bot_response))
     return chat_history