update app.py

GroundingDINO/groundingdino/version.py

@@ -1 +1 @@
-__version__ = "0.1.0"
+__version__ = '0.1.0'

app.py

@@ -1,12 +1,15 @@
 
-import subprocess, os, sys
+import subprocess, os, sys, time
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
-sys.path.insert(0, './GroundingDINO')
+result = subprocess.run(['pip', 'install', '-e', 'GroundingDINO'], check=True)
+print(f'pip install GroundingDINO = {result}')
 
 result = subprocess.run(['pip', 'list'], check=True)
-print(f'pip list result = {result}')
+print(f'pip list = {result}')
+
+sys.path.insert(0, './GroundingDINO')
 
 if not os.path.exists('./sam_vit_h_4b8939.pth'):
     result = subprocess.run(['wget', 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth'], check=True)
@@ -210,12 +213,22 @@ def run_grounded_sam(image_path, text_prompt, task_type, inpaint_prompt, box_thr
     # load image
     image_pil, image = load_image(image_path.convert("RGB"))
 
+    file_temp = int(time.time())
+
     # visualize raw image
-    image_pil.save(os.path.join(output_dir, "raw_image.jpg"))
+    # image_pil.save(os.path.join(output_dir, f"raw_image_{file_temp}.jpg"))
 
     # run grounding dino model
+    groundingdino_device = 'cpu'
+    if device != 'cpu':
+        try:
+            from groundingdino import _C
+            groundingdino_device = 'cuda:0'
+        except:
+            warnings.warn("Failed to load custom C++ ops. Running on CPU mode Only in groundingdino!")
+
     boxes_filt, pred_phrases = get_grounding_output(
-        groundingdino_model, image, text_prompt, box_threshold, text_threshold, device='cpu'
+        groundingdino_model, image, text_prompt, box_threshold, text_threshold, device=groundingdino_device
     )
 
     size = image_pil.size
@@ -250,9 +263,10 @@ def run_grounded_sam(image_path, text_prompt, task_type, inpaint_prompt, box_thr
         }
         # import ipdb; ipdb.set_trace()
         image_with_box = plot_boxes_to_image(image_pil, pred_dict)[0]
-        image_path = os.path.join(output_dir, "grounding_dino_output.jpg")
+        image_path = os.path.join(output_dir, f"grounding_dino_output_{file_temp}.jpg")
         image_with_box.save(image_path)
         image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+        os.remove(image_path)
         return image_result
     elif task_type == 'segment':
         assert sam_checkpoint, 'sam_checkpoint is not found!'
@@ -265,9 +279,10 @@ def run_grounded_sam(image_path, text_prompt, task_type, inpaint_prompt, box_thr
         for box, label in zip(boxes_filt, pred_phrases):
             show_box(box.numpy(), plt.gca(), label)
         plt.axis('off')
-        image_path = os.path.join(output_dir, "grounding_dino_output.jpg")
+        image_path = os.path.join(output_dir, f"grounding_seg_output_{file_temp}.jpg")
         plt.savefig(image_path, bbox_inches="tight")
         image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+        os.remove(image_path)
         return image_result
     elif task_type == 'inpainting':
         assert inpaint_prompt, 'inpaint_prompt is not found!'
@@ -277,9 +292,10 @@ def run_grounded_sam(image_path, text_prompt, task_type, inpaint_prompt, box_thr
         image_pil = Image.fromarray(image)
 
         image = sd_pipe(prompt=inpaint_prompt, image=image_pil, mask_image=mask_pil).images[0]
-        image_path = os.path.join(output_dir, "grounded_sam_inpainting_output.jpg")
+        image_path = os.path.join(output_dir, f"grounded_sam_inpainting_output_{file_temp}.jpg")
         image.save(image_path)
         image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
+        os.remove(image_path)
         return image_result
     else:
         print("task_type:{} error!".format(task_type))