adding changes for detection output

QA_bot.py

@@ -14,7 +14,6 @@ def display_mask_image(image_path):
         image = Image.open(image_path)
         st.image(image, caption='Final Mask', use_column_width=True)
     
-
 def tyre_synap_bot(filter_agent,image_file_path):
     if "messages" not in st.session_state:
         st.session_state.messages = []

extract_tools.py

@@ -9,7 +9,7 @@ from langchain_core.tools import tool,Tool
 from langchain_community.tools import DuckDuckGoSearchResults
 from langchain_groq import ChatGroq
 from utils import draw_panoptic_segmentation
-
+from typing import List 
 from tool_utils.clip_segmentation import CLIPSEG
 from tool_utils.yolo_world import YoloWorld
 from tool_utils.image_qualitycheck import brightness_check,gaussian_noise_check,snr_check
@@ -120,6 +120,7 @@ def clipsegmentation_mask(input_data:str)->str:
 @tool
 def generate_bounding_box_tool(input_data:str)->str:
     "use this tool when its is required to detect object and provide bounding boxes for the given image and list of objects"
+    print(input_data)
     data = input_data.split(",")
     image_path = data[0]
     object_prompts = data[1:]
@@ -201,10 +202,13 @@ def get_all_tools():
         bounding_box_generator = Tool(
             name = 'Bounding Box Generator',
             func = generate_bounding_box_tool,
-            description= "The tool helps to provide bounding boxes for the given image and list of objects\
-                .Use this tool when user ask to provide bounding boxes for the objects.if user has not specified the names of the objects \
-                then use the object extraction tool to identify the objects and then use this tool to generate the bounding boxes for the objects.\
-                The input to this tool is the path of the image and list of objects for which bounding boxes are to be generated"
+            description= """The tool helps to provide bounding boxes for the given image and list of objects
+                .Use this tool when user ask to provide bounding boxes for the objects.if user has not specified the names of the objects 
+                then use the object extraction tool to identify the objects and then use this tool to generate the bounding boxes for the objects.
+                The input to this tool is the path of the image and list of objects for which bounding boxes are to be generated
+                For Example :
+                    "action_input ": "image_store/<image_path>,person,dog,sand,"
+                """
         )
 
         object_extractor = Tool(

tool_utils/yolo_world.py

@@ -6,7 +6,7 @@ from typing import List
 import torch 
 import random 
 from ultralytics import YOLOWorld
-
+import json 
 class YoloWorld:
     def __init__(self,model_name = "yolov8x-worldv2.pt"):
         self.model = YOLOWorld(model_name)
@@ -57,7 +57,11 @@ class YoloWorld:
             cv2.putText(rgb_frame_copy, str(label), (c1[0], c1[1] - 2), 0, tl / 3, color_dict[label], thickness=tf, lineType=cv2.LINE_AA)
         return rgb_frame_copy
 
-
+    def format_detections(self,boxes,labels):
+        text = ""
+        for box ,label in zip(boxes,labels):
+            text+="{}\tBounding Box :{}\n".format(label,box)
+        return (text)
     def run_yolo_infer(self,image_path:str,object_prompts:List):
         processed_predictions = []
         bounding_boxes = []
@@ -73,20 +77,24 @@ class YoloWorld:
                 labels.append(result.names[int(box.cls.cpu())])
                 scores.append(round(float(box.conf.cpu()),2))
         
-        processed_predictions.append(dict(boxes= torch.tensor(bounding_boxes),
+        processed_predictions.append(dict(
+                                        boxes=  torch.tensor(bounding_boxes),
                                         labels= labels,
-                                        scores=torch.tensor(scores))
-                                    )
+                                        scores= torch.tensor(scores)
+                                        ) 
+                                    )                        
+       
         detected_image = self.draw_bboxes(rgb_frame=image_path,
                          boxes=processed_predictions[0]['boxes'],
                          labels=processed_predictions[0]['labels']
                          )
-        
+        predicted_data = self.format_detections(bounding_boxes,labels)
+        # save image 
         cv2.imwrite('final_mask.png', cv2.cvtColor(detected_image,cv2.COLOR_BGR2RGB))
-        return "Predicted image : final_mask.png . Details :{}".format(processed_predictions[0])
+        return "Predicted image : final_mask.png . \nDetails :\n{}".format(predicted_data)
                 
-# if __name__ == "__main__":
-#     yolo = YoloWorld()
-#     predicted_data = yolo.run_yolo_infer('../image_store/demo2.jpg',['person','hat','building'])
-#     print(predicted_data)
+if __name__ == "__main__":
+    yolo = YoloWorld()
+    predicted_data = yolo.run_yolo_infer('../image_store/demo2.jpg',['person','hat','building'])
+    print(predicted_data)