Update app.py

app.py

@@ -161,11 +161,33 @@ def video_vision(video_input_path, prompt, video_interval):
             prediction = translate_to_korean(prediction)
 
     if '[SEG]' in prediction and Visualizer is not None:
-        # 나머지 코드는 동일...
-        
-    return prediction, output_video if '[SEG]' in prediction else (prediction, None)
-    
-
+        _seg_idx = 0
+        pred_masks = result['prediction_masks'][_seg_idx]
+        seg_frames = []
+        for frame_idx in range(len(vid_frames)):
+            pred_mask = pred_masks[frame_idx]
+            temp_dir = tempfile.mkdtemp()
+            os.makedirs(temp_dir, exist_ok=True)
+            seg_frame = visualize(pred_mask, image_paths[frame_idx], temp_dir)
+            seg_frames.append(seg_frame)
+
+        output_video = "output_video.mp4"
+        frame = cv2.imread(seg_frames[0])
+        height, width, layers = frame.shape
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        video = cv2.VideoWriter(output_video, fourcc, new_fps, (width, height))
+
+        for img_path in seg_frames:
+            frame = cv2.imread(img_path)
+            video.write(frame)
+
+        video.release()
+        print(f"Video created successfully at {output_video}")
+
+        return prediction, output_video
+            
+    else:
+        return prediction, None
 
 # Gradio UI