fffiloni commited on
Commit
b01e53a
·
verified ·
1 Parent(s): 73bd08a

track tqdm in main function too

Browse files
Files changed (1) hide show
  1. app.py +10 -10
app.py CHANGED
@@ -301,9 +301,9 @@ def infer(video_in, trim_value, prompt, background_prompt, progress=gr.Progress(
301
  print(prompt)
302
  break_vid = get_frames(video_in)
303
 
304
- frames_list= break_vid[0]
305
  fps = break_vid[1]
306
- n_frame = int(trim_value*fps)
307
 
308
  if n_frame >= len(frames_list):
309
  print("video is shorter than the cut value")
@@ -313,36 +313,36 @@ def infer(video_in, trim_value, prompt, background_prompt, progress=gr.Progress(
313
  with_green_result_frames = []
314
  with_matte_result_frames = []
315
 
316
- print("set stop frames to: " + str(n_frame))
317
  bg_already = False
318
- for i in frames_list[0:int(n_frame)]:
 
 
319
  to_numpy_i = Image.open(i).convert("RGB")
320
- #need to convert to numpy
321
- # Convert the image to a NumPy array
322
  image_array = np.array(to_numpy_i)
323
 
324
  results = run_grounded_sam(image_array, prompt, "text", background_prompt, bg_already)
325
  bg_already = True
 
326
  bg_img = Image.fromarray(results[0])
327
  green_img = Image.fromarray(results[1])
328
  matte_img = Image.fromarray(results[2])
329
 
330
-
331
  # exporting the images
332
  bg_img.save(f"bg_result_img-{i}.jpg")
333
  with_bg_result_frames.append(f"bg_result_img-{i}.jpg")
 
334
  green_img.save(f"green_result_img-{i}.jpg")
335
  with_green_result_frames.append(f"green_result_img-{i}.jpg")
 
336
  matte_img.save(f"matte_result_img-{i}.jpg")
337
  with_matte_result_frames.append(f"matte_result_img-{i}.jpg")
338
- print("frame " + i + "/" + str(n_frame) + ": done;")
339
 
340
  vid_bg = create_video(with_bg_result_frames, fps, "bg")
341
  vid_green = create_video(with_green_result_frames, fps, "greenscreen")
342
  vid_matte = create_video(with_matte_result_frames, fps, "matte")
343
 
344
- bg_already = False
345
- print("finished !")
346
 
347
  return vid_bg, vid_green, vid_matte
348
 
 
301
  print(prompt)
302
  break_vid = get_frames(video_in)
303
 
304
+ frames_list = break_vid[0]
305
  fps = break_vid[1]
306
+ n_frame = int(trim_value * fps)
307
 
308
  if n_frame >= len(frames_list):
309
  print("video is shorter than the cut value")
 
313
  with_green_result_frames = []
314
  with_matte_result_frames = []
315
 
316
+ print("set stop frames to:", n_frame)
317
  bg_already = False
318
+
319
+ # tqdm hooked into gr.Progress
320
+ for i in tqdm(frames_list[:n_frame], desc="Inferring frames", unit="frame"):
321
  to_numpy_i = Image.open(i).convert("RGB")
 
 
322
  image_array = np.array(to_numpy_i)
323
 
324
  results = run_grounded_sam(image_array, prompt, "text", background_prompt, bg_already)
325
  bg_already = True
326
+
327
  bg_img = Image.fromarray(results[0])
328
  green_img = Image.fromarray(results[1])
329
  matte_img = Image.fromarray(results[2])
330
 
 
331
  # exporting the images
332
  bg_img.save(f"bg_result_img-{i}.jpg")
333
  with_bg_result_frames.append(f"bg_result_img-{i}.jpg")
334
+
335
  green_img.save(f"green_result_img-{i}.jpg")
336
  with_green_result_frames.append(f"green_result_img-{i}.jpg")
337
+
338
  matte_img.save(f"matte_result_img-{i}.jpg")
339
  with_matte_result_frames.append(f"matte_result_img-{i}.jpg")
 
340
 
341
  vid_bg = create_video(with_bg_result_frames, fps, "bg")
342
  vid_green = create_video(with_green_result_frames, fps, "greenscreen")
343
  vid_matte = create_video(with_matte_result_frames, fps, "matte")
344
 
345
+ print("finished!")
 
346
 
347
  return vid_bg, vid_green, vid_matte
348