Update app.py

app.py

@@ -21,7 +21,7 @@ import math
 from scipy.optimize import linear_sum_assignment
 from obj365_name import categories as OBJ365_CATEGORIESV2
 import copy
-
+import skvideo.io
 
 this_dir = path.dirname(path.abspath(__file__))
 
@@ -616,10 +616,14 @@ def segment_video(video, prompt_mode, categoryname, custom_category, expressiong
         size = (ori_width,ori_height)    
         print('writing video...')
         output_file = "test.mp4"
-        out = cv2.VideoWriter(output_file,cv2.VideoWriter_fourcc(*'avc1'), read_fps, size) 
+        writer = skvideo.io.FFmpegWriter(output_file)
         for i in range(len(ourput_frames)):
-            out.write(ourput_frames[i])
-        out.release()
+            writer.writeFrame(ourput_frames[i])
+        writer.close()
+        # out = cv2.VideoWriter(output_file,cv2.VideoWriter_fourcc(*'avc1'), read_fps, size) 
+        # for i in range(len(ourput_frames)):
+        #     out.write(ourput_frames[i])
+        # out.release()
         del out_logits, out_masks, out_embds, out_boxes, pred_masks
         torch.cuda.empty_cache()
         return output_file