Update app.py

app.py

@@ -1,13 +1,13 @@
 import os
 import gradio as gr
-from transformers import DetrForObjectDetection, DetrImageProcessor
-from langchain_google_genai.chat_models import ChatGoogleGenerativeAI  # For Chat Gemini
+from transformers import DetrImageProcessor, DetrForObjectDetection
+from langchain_google_genai.chat_models import ChatGoogleGenerativeAI  # Import Gemini
 from PIL import Image
 import torch
 import json
 import requests
 
-# Load credentials (stringified JSON) from environment variable
+# Load credentials (stringified JSON) from environment variable for Gemini
 credentials_string = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
 if not credentials_string:
     raise ValueError("GOOGLE_APPLICATION_CREDENTIALS is not set in the environment!")
@@ -22,23 +22,37 @@ with open("service_account.json", "w") as f:
 # Set the GOOGLE_APPLICATION_CREDENTIALS environment variable to the temporary file
 os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "service_account.json"
 
-# Initialize Chat Gemini model
+# Initialize Gemini model (chatbot)
 llm = ChatGoogleGenerativeAI(model='gemini-1.5-pro')
 
-# Load the model and processor for DETR (for object detection)
+# Initialize DETR model and processor for object detection
 processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
 model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 
-# Global chat history
+# Load COCO class labels (from the official COCO dataset)
+COCO_CLASSES = [
+    'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 
+    'traffic light', 'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 
+    'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 
+    'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 
+    'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 
+    'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 
+    'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 
+    'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 
+    'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush'
+]
+
+# Global chat history variable
 chat_history = []
 
+# Function for chatting with Gemini
 def chat_with_gemini(message):
     global chat_history
-    # Get a response from the Gemini model
     bot_response = llm.predict(message)  # This will interact with the Gemini model
     chat_history.append((message, bot_response))
     return chat_history
 
+# Function for analyzing the uploaded image
 def analyze_image(image_path):
     global chat_history
     try:
@@ -73,7 +87,6 @@ def analyze_image(image_path):
         chat_history.append(("Error during image analysis", error_msg))
         return chat_history
 
-
 # Build the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Ken Chatbot")
@@ -128,5 +141,5 @@ with gr.Blocks() as demo:
     </style>
     """)
 
-# Launch for Hugging Face Spaces
+# Launch the Gradio interface
 demo.launch()