error fix

app.py

@@ -23,24 +23,35 @@ def respond(input_image):
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
         cv2.imwrite(temp_file.name, input_image)
 
-        DINOGPT = CustomDetectionModel(
-            detection_model=GroundingDINO(
-                CaptionOntology({dino_prompt: dino_prompt})
-            ),
-            classification_model=GPT4V(
-                CaptionOntology({k: k for k in gpt_prompt.split(", ")}),
-                api_key=api_key  # Use the hardcoded API key
+        try:
+            DINOGPT = CustomDetectionModel(
+                detection_model=GroundingDINO(
+                    CaptionOntology({dino_prompt: dino_prompt})
+                ),
+                classification_model=GPT4V(
+                    CaptionOntology({k: k for k in gpt_prompt.split(", ")}),
+                    api_key=api_key  # Use the hardcoded API key
+                )
             )
-        )
 
-        results = DINOGPT.predict(temp_file.name)
+            results = DINOGPT.predict(temp_file.name)
 
-        result_image = plot(
-            image=cv2.imread(temp_file.name),
-            detections=results,
-            classes=gpt_prompt.split(", "),
-            raw=True
-        )
+            # Assuming the results are a list of dictionaries with 'class' keys
+            # Extracting class names from results
+            predicted_classes = [result['class'] for result in results]
+            # Filtering the ontology prompts to include only those that match the predicted classes
+            filtered_ontology = {k: v for k, v in DINOGPT.classification_model.ontology.prompts().items() if k in predicted_classes}
+
+            result_image = plot(
+                image=cv2.imread(temp_file.name),
+                detections=results,
+                classes=list(filtered_ontology.keys()),  # Use the filtered ontology
+                raw=True
+            )
+        except ValueError as e:
+            print("An error occurred:", e)
+            # Handle the error appropriately, perhaps by returning an error message or default image
+            result_image = None  # Replace with your error handling logic
 
     return result_image