fffiloni commited on
Commit
fe10180
1 Parent(s): 19ed362

added tab for auto process

Browse files
Files changed (1) hide show
  1. app.py +172 -5
app.py CHANGED
@@ -1,8 +1,14 @@
1
  import os
2
  import shutil
3
  import subprocess
 
 
 
 
4
  import gradio as gr
5
 
 
 
6
  # execute a CLI command
7
  def execute_command(command: str) -> None:
8
  subprocess.run(command, check=True)
@@ -92,6 +98,157 @@ def infer(video_frames, masks_frames, project_name):
92
 
93
  return "done", results_folder_content[0], results_folder_content[1]
94
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
95
  css="""
96
  #col-container{
97
  margin: 0 auto;
@@ -111,12 +268,21 @@ with gr.Blocks(css=css) as demo:
111
  """)
112
 
113
  with gr.Row():
114
- with gr.Column():
115
- project_name = gr.Textbox(label="Name your project", value="my-new-project")
116
- video_frames = gr.File(label="Video frames", file_types=["image"], file_count="multiple")
117
- masks_frames = gr.File(label="Masks frames", file_types=["image"], file_count="multiple")
 
118
 
119
- submit_btn = gr.Button("Submit")
 
 
 
 
 
 
 
 
120
 
121
  with gr.Column():
122
  result = gr.Textbox(label="Result")
@@ -126,5 +292,6 @@ with gr.Blocks(css=css) as demo:
126
 
127
 
128
  submit_btn.click(fn=infer, inputs=[video_frames, masks_frames, project_name], outputs=[result, res_masked, res_files])
 
129
 
130
  demo.queue(max_size=12).launch()
 
1
  import os
2
  import shutil
3
  import subprocess
4
+ import cv2
5
+ from PIL import Image
6
+ from moviepy.editor import *
7
+ from gradio_client import Client
8
  import gradio as gr
9
 
10
+ matte_client = Client("https://fffiloni-video-matting-anything.hf.space/")
11
+
12
  # execute a CLI command
13
  def execute_command(command: str) -> None:
14
  subprocess.run(command, check=True)
 
98
 
99
  return "done", results_folder_content[0], results_folder_content[1]
100
 
101
+
102
+ def get_frames(video_in, img_type):
103
+ frames = []
104
+ #resize the video
105
+ clip = VideoFileClip(video_in)
106
+
107
+ #check fps
108
+ if clip.fps > 30:
109
+ print("vide rate is over 30, resetting to 30")
110
+ clip_resized = clip.resize(height=512)
111
+ clip_resized.write_videofile(f"{img_type}_video_resized.mp4", fps=30)
112
+ else:
113
+ print("video rate is OK")
114
+ clip_resized = clip.resize(height=512)
115
+ clip_resized.write_videofile(f"{img_type}_video_resized.mp4", fps=clip.fps)
116
+
117
+ print("video resized to 512 height")
118
+
119
+ # Opens the Video file with CV2
120
+ cap= cv2.VideoCapture(f"{img_type}_video_resized.mp4")
121
+
122
+ fps = cap.get(cv2.CAP_PROP_FPS)
123
+ print("video fps: " + str(fps))
124
+ i=0
125
+ while(cap.isOpened()):
126
+ ret, frame = cap.read()
127
+ if ret == False:
128
+ break
129
+ if img_type == "source":
130
+ cv2.imwrite('kang'+str(i)+'.jpg',frame)
131
+ frames.append('kang'+str(i)+'.jpg')
132
+ elif img_type == "mask":
133
+ cv2.imwrite('kang'+str(i)+'.png',frame)
134
+ frames.append('kang'+str(i)+'.png')
135
+ i+=1
136
+
137
+ cap.release()
138
+ cv2.destroyAllWindows()
139
+ print("broke the video into frames")
140
+
141
+ return frames, fps
142
+
143
+ def get_matte(video_in, subject_to_remove):
144
+ print("Trying to call video matting")
145
+ result = matte_client.predict(
146
+ f"{video_in}", # str (filepath on your computer (or URL) of file) in 'parameter_4' Video component
147
+ 10, # int | float (numeric value between 0 and 10) in 'Cut video at (s)' Slider component
148
+ f"{subject_to_remove}", # str in 'Text prompt' Textbox component
149
+ "", # str in 'Background prompt' Textbox component
150
+ api_name="/go_matte"
151
+ )
152
+ print(result)
153
+
154
+ return result[2]
155
+
156
+
157
+ def infer_auto(project_name, video_in, subject_to_remove):
158
+
159
+ print(video_in)
160
+ matte_video = get_matte(video_in, subject_to_remove)
161
+
162
+ video_frames = get_frames(video_in, "source")
163
+ print(video_frames[0])
164
+ #video_frames_paths = [os.path.abspath(image) for image in video_frames[0]]
165
+
166
+
167
+ masks_frames = get_frames(matte_video, "mask")
168
+ print(masks_frames[0])
169
+ #masks_frames_paths = [os.path.abspath(image) for image in masks_frames[0]]
170
+
171
+ # Create the directory if it doesn't exist
172
+ my_video_directory = f"{project_name}"
173
+ if not os.path.exists(my_video_directory):
174
+ os.makedirs(my_video_directory)
175
+
176
+ # Assuming 'images' is a list of image file paths
177
+ for idx, image in enumerate(video_frames[0]):
178
+ # Get the base file name (without path) from the original location
179
+ image_name = os.path.basename(image)
180
+
181
+ # Construct the destination path in the working directory
182
+ destination_path = os.path.join(my_video_directory, image_name)
183
+
184
+ # Copy the image from the original location to the working directory
185
+ shutil.copy(image, destination_path)
186
+
187
+ # Print the image name and its corresponding save path
188
+ print(f"Image {idx + 1}: {image_name} copied to {destination_path}")
189
+
190
+
191
+ # Create the directory if it doesn't exist
192
+ my_masks_directory = f"{project_name}_masks"
193
+ if not os.path.exists(my_masks_directory):
194
+ os.makedirs(my_masks_directory)
195
+
196
+ # Assuming 'images' is a list of image file paths
197
+ for idx, image in enumerate(masks_frames[0]):
198
+ # Get the base file name (without path) from the original location
199
+ image_name = os.path.basename(image)
200
+
201
+ # Construct the destination path in the working directory
202
+ destination_path = os.path.join(my_masks_directory, image_name)
203
+
204
+ # Copy the image from the original location to the working directory
205
+ shutil.copy(image, destination_path)
206
+
207
+ # Print the image name and its corresponding save path
208
+ print(f"Image {idx + 1}: {image_name} copied to {destination_path}")
209
+
210
+ #video_frames_folder = "inputs/object_removal/bmx-trees"
211
+ #masks_folder = "inputs/object_removal/bmx-trees_mask"
212
+
213
+ video_frames_folder = f"{my_video_directory}"
214
+ masks_folder = f"{my_masks_directory}"
215
+
216
+ # Create the "results" folder if it doesn't exist
217
+ output_folder = "results"
218
+ if not os.path.exists(output_folder):
219
+ os.makedirs(output_folder)
220
+
221
+
222
+ #bmx_trees_folder = os.path.join(output_folder, "bmx-trees")
223
+
224
+ command = [
225
+ f"python",
226
+ f"inference_propainter.py",
227
+ f"--video={video_frames_folder}",
228
+ f"--mask={masks_folder}",
229
+ f"--output={output_folder}"
230
+ ]
231
+
232
+ execute_command(command)
233
+
234
+ # Get the list of files in the "results" folder
235
+ result_files = os.listdir(output_folder)
236
+
237
+ # Print the content of the "results" folder
238
+ print(f"Contents of the {output_folder} folder:")
239
+ for item in result_files:
240
+ print(item)
241
+
242
+ # List the content of the "bmx-trees" folder within "results"
243
+ results_folder = os.path.join(output_folder, f"{project_name}")
244
+ results_folder_content = [os.path.join(results_folder, item) for item in os.listdir(results_folder)]
245
+
246
+ print(f"Contents of the {results_folder} folder:")
247
+ for item in results_folder_content:
248
+ print(item)
249
+
250
+ return "done", results_folder_content[0], results_folder_content[1]
251
+
252
  css="""
253
  #col-container{
254
  margin: 0 auto;
 
268
  """)
269
 
270
  with gr.Row():
271
+ with gr.Tab("Manual"):
272
+ with gr.Column():
273
+ project_name = gr.Textbox(label="Name your project", value="my-new-project")
274
+ video_frames = gr.File(label="Video frames", file_types=["image"], file_count="multiple")
275
+ masks_frames = gr.File(label="Masks frames", file_types=["image"], file_count="multiple")
276
 
277
+ submit_btn = gr.Button("Submit")
278
+
279
+ with gr.Tab("Auto"):
280
+ with gr.Column():
281
+ project_name_2 = gr.Textbox(label="Name your project", value="my-new-project")
282
+ video_in = gr.Video(label="Source video", source="upload", format="mp4")
283
+ subject_to_remove = gr.Textbox(label="Subject to remove")
284
+
285
+ submit_auto_btn = gr.Button("Submit")
286
 
287
  with gr.Column():
288
  result = gr.Textbox(label="Result")
 
292
 
293
 
294
  submit_btn.click(fn=infer, inputs=[video_frames, masks_frames, project_name], outputs=[result, res_masked, res_files])
295
+ submit_auto_btn.click(fn=infer_auto, inputs=[project_name, video_in, subject_to_remove], outputs=[result, res_masked, res_files])
296
 
297
  demo.queue(max_size=12).launch()