Update modules/masking_module.py

modules/masking_module.py

@@ -9,6 +9,7 @@ from typing import Any
 import supervision as sv
 from sam2.build_sam import build_sam2, build_sam2_video_predictor
 from sam2.sam2_image_predictor import SAM2ImagePredictor
+import time
 
 device = torch.device('cuda')
 
@@ -134,7 +135,8 @@ def masking_process(image,obj):
     # task_prompt = '<REGION_TO_SEGMENTATION>'
     # # task_prompt = '<OPEN_VOCABULARY_DETECTION>'
     # print(type(task_prompt),type(obj))
-    print('1')
+    # print('1')
+    start_time = time.time()
     image = Image.fromarray(image).convert("RGB")
 
     # results = florence2(image,task_prompt, text_input=obj)
@@ -150,17 +152,18 @@ def masking_process(image,obj):
     # obj = "Tiger"
 
     Florence_results = florence2(image,task_prompt, text_input=obj)
-    print('2')
+    # print('2')
     SAM_IMAGE_MODEL = load_sam_image_model(device=device)
-    print('3')
+    # print('3')
     detections = sv.Detections.from_lmm(
                 lmm=sv.LMM.FLORENCE_2,
                 result=Florence_results,
                 resolution_wh=image.size
             )
-    print('4')
+    # print('4')
     response = run_sam_inference(SAM_IMAGE_MODEL, image, detections)
-    print('5')
+    print(f'Time taken by masking model: {time.time() - start}')
+    # print('5')
     if response['code'] == 400:
         print("no object found")
         return "no object found"