Update app.py

app.py

@@ -15,6 +15,7 @@ from vincenty import vincenty
 import json
 from collections import Counter
 import mediapy
+import ffmpy
 
 #from depth_anything.dpt import DepthAnything
 #from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
@@ -58,6 +59,15 @@ def zip_files(files_in, files_out):
             zipObj.write(file, file.split("/")[-1])
     return "depth_result.zip"
 
+def create_alpha_video(frames, fps, type):
+    print("building composite video result")
+    imgs = []
+    for j, img in enumerate(frames):
+        imgs.append(cv2.cvtColor(cv2.imread(img).astype(np.uint8), cv2.COLOR_RGBA2RGB))
+
+    mediapy.write_video(type + "_result.mp4", imgs, fps=fps)
+    return type + "_result.mp4"
+
 def create_video(frames, fps, type):
     print("building video result")
     imgs = []
@@ -159,6 +169,7 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data
         depth_frames = []
         orig_frames = []
         backup_frames = []
+        rgba_frames = []
         thumbnail_old = []
 
         while raw_video.isOpened():
@@ -222,15 +233,11 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data
               thumbnail_old = thumbnail
 
             blur_frame = blur_image(raw_frame, depth_color, blur_data)
+            rgba_frame = cv2.cvtColor(cv2.cvtColor(blur_frame, cv2.COLOR_BGR2RGB), cv2.COLOR_RGB2RGBA)
+            rgba_frame[:,:,3] = depth_gray
 
-            # encoding depth within original video
-            blur_frame = (int(blur_frame / 17) * 17).astype(np.uint8)
-            depth_r = int(depth_gray / 17).astype(np.uint8)
-            # may use green channel for 16 levels of opacity
-            depth_b = depth_gray - depth_r * 17
-            blur_frame[:,:,0] = blur_frame[:,:,0] + depth_r
-            # blur_frame[:,:,1] = blur_frame[:,:,1] + opacity_g
-            blur_frame[:,:,2] = blur_frame[:,:,2] + depth_b
+            cv2.imwrite(f"f{count}.png", rgba_frame)
+            rgba_frames.append(f"f{count}.png")
             
             cv2.imwrite(f"f{count}.jpg", blur_frame)
             orig_frames.append(f"f{count}.jpg")
@@ -250,6 +257,7 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data
 
         final_vid = create_video(orig_frames, frame_rate, "orig")
         depth_vid = create_video(depth_frames, frame_rate, "depth")
+        alpha_vid = create_alpha_video(rgba_frames, frame_rate, "rgba")
             
         final_zip = zip_files(orig_frames, depth_frames)
         raw_video.release()
@@ -272,7 +280,7 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data
         else:
             gradient = cv2.imread('./gradient_small.png').astype(np.uint8)
         
-        return final_vid, final_zip, frames, masks[frame_selected], depths, depth_vid #output_path
+        return alpha_vid, final_zip, frames, masks[frame_selected], depths, depth_vid #output_path
 
 def depth_edges_mask(depth):
     """Returns a mask of edges in the depth map.