Update app.py

app.py

@@ -114,7 +114,7 @@ def create_video(frames, fps, type):
     return f"video_{type}_result.mp4"
 
 
-def run_grounded_sam(input_image, text_prompt, task_type, background_prompt):
+def run_grounded_sam(input_image, text_prompt, task_type, background_prompt, bg_already):
     background_type = "generated_by_text"
     box_threshold = 0.25
     text_threshold = 0.25
@@ -262,6 +262,7 @@ def run_grounded_sam(input_image, text_prompt, task_type, background_prompt):
     ### alpha matte
     alpha_rgb = cv2.cvtColor(np.uint8(alpha_pred*255), cv2.COLOR_GRAY2RGB)
     ### com img with background
+    global background_img
     if background_type == 'real_world_sample':
         background_img_file = os.path.join('assets/backgrounds', random.choice(background_list))
         background_img = cv2.imread(background_img_file)
@@ -273,7 +274,9 @@ def run_grounded_sam(input_image, text_prompt, task_type, background_prompt):
         if background_prompt is None:
             print('Please input non-empty background prompt')
         else:
-            background_img = generator(background_prompt).images[0]
+            if bg_already is False:
+                background_img = generator(background_prompt).images[0]
+            else: 
             background_img = np.array(background_img)
             background_img = cv2.resize(background_img, (image_ori.shape[1], image_ori.shape[0]))
             com_img = alpha_pred[..., None] * image_ori + (1 - alpha_pred[..., None]) * np.uint8(background_img)
@@ -301,15 +304,15 @@ def infer(video_in, trim_value, prompt, background_prompt):
     with_matte_result_frames = []
     
     print("set stop frames to: " + str(n_frame))
-    
+    bg_already = False
     for i in frames_list[0:int(n_frame)]:
         to_numpy_i = Image.open(i).convert("RGB")
         #need to convert to numpy
         # Convert the image to a NumPy array
         image_array = np.array(to_numpy_i)
 
-        results = run_grounded_sam(image_array, prompt, "text", background_prompt)
-
+        results = run_grounded_sam(image_array, prompt, "text", background_prompt, bg_already)
+        bg_already = True
         bg_img = Image.fromarray(results[0])
         green_img = Image.fromarray(results[1])
         matte_img = Image.fromarray(results[2])
@@ -327,6 +330,8 @@ def infer(video_in, trim_value, prompt, background_prompt):
     vid_bg = create_video(with_bg_result_frames, fps, "bg")
     vid_green = create_video(with_green_result_frames, fps, "greenscreen")
     vid_matte = create_video(with_matte_result_frames, fps, "matte")
+
+    bg_already = False
     print("finished !")
     
     return vid_bg, vid_green, vid_matte