added maskformer based object extraction

extract_tools.py

@@ -124,39 +124,25 @@ def generate_bounding_box_tool(input_data:str)->str:
     object_data = yolo_world_model.run_yolo_infer(image_path,object_prompts)
     return object_data
 
+
 @tool 
-def object_extraction(img_path:str)->str:
+def object_extraction(image_path:str)->str:
     "Use this tool to identify the objects within the image"
-    
-    hf_model = "Salesforce/blip-image-captioning-base"
-    if img_path.startswith('https'):
-        image = Image.open(requests.get(img_path, stream=True).raw).convert('RGB')
-    else:
-        image = Image.open(img_path).convert('RGB')
-    try:
-        processor = BlipProcessor.from_pretrained(hf_model)
-        caption_model = BlipForConditionalGeneration.from_pretrained(hf_model).to(device)
-    except:
-        logging.error("unable to load the Blip model ")
-    
-    logging.info("Image Caption model loaded ! ")
-    
-    # unconditional image captioning
-    inputs = processor(image, return_tensors ='pt').to(device)
-    output = caption_model.generate(**inputs, max_new_tokens=50)
-    llm = get_groq_model() 
-    getobject_chain = create_object_extraction_chain(llm=llm)
-    
-    extracted_objects = getobject_chain.invoke({
-        'context': processor.decode(output[0], skip_special_tokens=True)
-    }).objects
-
-    print("Extracted objects : ",extracted_objects)
-    ## clear the GPU cache 
+    objects = []
+    maskformer_model.to(device)
+    image = cv2.imread(image_path)
+    image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
+    inputs = maskformer_processor(image, return_tensors="pt")
+    inputs.to(device)
     with torch.no_grad():
-        torch.cuda.empty_cache()
-    
-    return extracted_objects.split(',') 
+        outputs = maskformer_model(**inputs)
+    prediction = maskformer_processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.shape[:2]])[0]
+    segments_info = prediction['segments_info']
+    for segment in segments_info:
+        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)
 
 @tool 
 def get_image_quality(image_path:str)->str:

utils.py

@@ -50,4 +50,8 @@ def draw_bboxes(rgb_frame,boxes,labels,color=None,line_thickness=3):
         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)
-    return rgb_frame_copy
+    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")