added object detection imporved

extract_tools.py

@@ -11,7 +11,6 @@ from langchain_groq import ChatGroq
 from utils import draw_panoptic_segmentation
 
 from tool_utils.clip_segmentation import CLIPSEG
-from tool_utils.object_extractor import create_object_extraction_chain
 from tool_utils.yolo_world import YoloWorld
 from tool_utils.image_qualitycheck import brightness_check,gaussian_noise_check,snr_check
 
@@ -34,6 +33,10 @@ try:
 except:
     logging.error("Unable to Maskformer model {}".format(err))
 
+try:
+    yolo_world_model= YoloWorld()
+except : 
+    logging.error("Unable to Yolo world model {}".format(err))
 
 def get_groq_model(model_name = "gemma2-9b-it"):
     os.environ.get("GROQ_API_KEY")
@@ -117,7 +120,6 @@ 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"
-    yolo_world_model= YoloWorld()
     data = input_data.split(",")
     image_path = data[0]
     object_prompts = data[1:]
@@ -142,7 +144,8 @@ def object_extraction(image_path:str)->str:
         segment_label_id = segment['label_id']
         segment_label = maskformer_model.config.id2label[segment_label_id]
         objects.append(segment_label) 
-    return "Detected objects are:  "+ " ".join( objects)
+    
+    return "Detected objects are:  "+ ",".join( objects)
 
 @tool 
 def get_image_quality(image_path:str)->str:

tool_utils/yolo_world.py

@@ -32,14 +32,16 @@ class YoloWorld:
         return object_details
     
     @staticmethod
-    def draw_bboxes(rgb_frame,boxes,labels,color=None,line_thickness=3):
-        rgb_frame = cv2.imread(rgb_frame)
-        rgb_frame = cv2.cvtColor(rgb_frame,cv2.COLOR_BGR2RGB)
+    def draw_bboxes(rgb_frame,boxes,labels,line_thickness=3):
+        rgb_frame = cv2.cvtColor(cv2.imread(rgb_frame),cv2.COLOR_BGR2RGB)
         
         tl = line_thickness or round(0.002 * (rgb_frame.shape[0] + rgb_frame.shape[1]) / 2) + 1  # line/font thickness
         rgb_frame_copy = rgb_frame.copy()
-        if color is None :
-            color = color or [random.randint(0, 255) for _ in range(3)]
+        color_dict = {}
+        # color = color or [random.randint(0, 255) for _ in range(3)]
+        for item in np.unique(np.asarray(labels)):
+            color_dict[item] = [random.randint(28, 255) for _ in range(3)]
+        
         for box,label in zip(boxes,labels):
             if box.type() == 'torch.IntTensor':
                 box = box.numpy()
@@ -47,31 +49,32 @@ class YoloWorld:
             x1,y1,x2,y2 = box
             c1,c2  = (x1,y1),(x2,y2)
             # Draw rectangle
-            cv2.rectangle(rgb_frame_copy, c1,c2, color, thickness=tl, lineType=cv2.LINE_AA)
-            
+            cv2.rectangle(rgb_frame_copy, c1,c2, color_dict[label], thickness=tl, lineType=cv2.LINE_AA)
             tf = max(tl - 1, 1)  # font thickness
             # label = label2id[int(label.numpy())]
             t_size = cv2.getTextSize(str(label), 0, fontScale=tl / 3, thickness=tf)[0]
             c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
-            cv2.putText(rgb_frame_copy, str(label), (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+            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 run_yolo_infer(self,image_path:str,object_prompts:List):
-        self.model.set_classes(object_prompts)
-        results = self.model.predict(image_path)
         processed_predictions = []
         bounding_boxes = []
         labels = []
         scores = []
+
+        self.model.set_classes(object_prompts)
+        results = self.model.predict(image_path)
         for result in results:
             for i,box in enumerate(result.boxes):
                 x1, y1, x2, y2 = np.array(box.xyxy.cpu(), dtype=np.int32).squeeze()
                 bounding_boxes.append([x1,y1,x2,y2])
-                labels.append(int(box.cls.cpu()))
+                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),
-                                        labels= torch.IntTensor(labels),
+                                        labels= labels,
                                         scores=torch.tensor(scores))
                                     )
         detected_image = self.draw_bboxes(rgb_frame=image_path,
@@ -80,5 +83,10 @@ class YoloWorld:
                          )
         
         cv2.imwrite('final_mask.png', cv2.cvtColor(detected_image,cv2.COLOR_BGR2RGB))
-        return "Predicted image mask : final_mask.jpg . Details :{}".format(processed_predictions[0]['boxes'])
+        return "Predicted image : final_mask.jpg . Details :{}".format(processed_predictions[0])
                 
+if __name__ == "__main__":
+    yolo = YoloWorld()
+    predicted_data = yolo.run_yolo_infer('../image_store/demo2.jpg',['person','hat','building'])
+    print(predicted_data)
+

tool_utils/yolov8x-worldv2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41e771bfbbb8894dd857f3fef7cac3b3578dffd49fd3547101efa6a606a02a0e
+size 146355704

utils.py

@@ -1,3 +1,4 @@
+import numpy as np
 from collections import defaultdict
 import matplotlib.pyplot as plt
 import matplotlib.patches as mpatches
@@ -28,14 +29,17 @@ def draw_panoptic_segmentation(model,segmentation, segments_info):
     return 'final_mask.png'
 
 
-def draw_bboxes(rgb_frame,boxes,labels,color=None,line_thickness=3):
+def draw_bboxes(rgb_frame,boxes,labels,line_thickness=3):
     rgb_frame = cv2.imread(rgb_frame)
     # rgb_frame = cv2.cvtColor(rgb_frame,cv2.COLOR_BGR2RGB)
     
     tl = line_thickness or round(0.002 * (rgb_frame.shape[0] + rgb_frame.shape[1]) / 2) + 1  # line/font thickness
     rgb_frame_copy = rgb_frame.copy()
-    if color is None :
-        color = color or [random.randint(0, 255) for _ in range(3)]
+    color_dict = {}
+    # color = color or [random.randint(0, 255) for _ in range(3)]
+    for item in np.unique(np.asarray(labels)):
+        color_dict[item] = [random.randint(28, 255) for _ in range(3)]
+    
     for box,label in zip(boxes,labels):
         if box.type() == 'torch.IntTensor':
             box = box.numpy()
@@ -43,15 +47,10 @@ def draw_bboxes(rgb_frame,boxes,labels,color=None,line_thickness=3):
         x1,y1,x2,y2 = box
         c1,c2  = (x1,y1),(x2,y2)
         # Draw rectangle
-        cv2.rectangle(rgb_frame_copy, c1,c2, color, thickness=tl, lineType=cv2.LINE_AA)
-        
+        cv2.rectangle(rgb_frame_copy, c1,c2, color_dict[label], thickness=tl, lineType=cv2.LINE_AA)
         tf = max(tl - 1, 1)  # font thickness
         # label = label2id[int(label.numpy())]
         t_size = cv2.getTextSize(str(label), 0, fontScale=tl / 3, thickness=tf)[0]
         c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
-        cv2.putText(rgb_frame_copy, str(label), (c1[0], c1[1] - 2), 0, tl / 3, [225, 0, 255], thickness=tf, lineType=cv2.LINE_AA)
+        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 object_extraction_using_maskformer(image_path):
-    processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
-    model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")