File size: 5,570 Bytes
8b09391 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
import os
import warnings
import cv2
import streamlit as st
from PIL import Image, ImageDraw
import redirect as rd
from main import segment_video
warnings.filterwarnings("ignore")
def load_image(image_path):
return Image.open(image_path)
def extract_first_frame(video_path, output_image_path):
"""
Extract the first frame from a video file and save it to disk.
Parameters:
video_path (str): Path to the video file.
output_image_path (str): Path to save the extracted frame.
Returns:
str: Path to the saved frame.
"""
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
raise ValueError(f"Error: Unable to open video file: {video_path}")
ret, frame = cap.read()
cap.release()
if not ret:
raise ValueError("Error: Unable to read the first frame from the video.")
cv2.imwrite(output_image_path, frame)
return output_image_path
st.title("Video Background Removal")
st.write(
"This app uses the Mobile-SAM model to remove the background from a video. "
"The model is based on the paper [Faster Segment Anything: Towards Lightweight SAM for Mobile Applications](https://arxiv.org/abs/2306.14289)."
)
st.write(
"How to use: Upload a video and click 'Segment Video'. The app will then process the video and remove the background. "
"You can also use a bounding box to specify the area to segment. "
"The app will then output the segmented video, that you can download. "
"Do not hesitate to hit the 'Stop/Reset' button if you encounter any issues (it usually solves them all) or want to start over."
)
video_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov"])
if video_file is not None:
st.video(video_file)
with open("temp_video.mp4", "wb") as f:
f.write(video_file.getbuffer())
if not os.path.exists("./temp_images"):
os.makedirs("./temp_images")
frame_path = extract_first_frame("temp_video.mp4", "temp_frame.jpg")
use_bbox = st.checkbox("Use bounding box", value=False)
background_color = st.color_picker("Background keying color", "#009000")
initial_frame = load_image(frame_path)
original_width, original_height = initial_frame.width, initial_frame.height
if use_bbox:
col1, col2 = st.columns(2)
with col1:
xmin = st.slider("xmin", 0, original_width, original_width // 4)
ymin = st.slider("ymin", 0, original_height, original_height // 4)
with col2:
xmax = st.slider("xmax", 0, original_width, original_width // 2)
ymax = st.slider("ymax", 0, original_height, original_height // 2)
draw = ImageDraw.Draw(initial_frame)
draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=3)
st.image(initial_frame, caption="Bounding Box Preview", use_column_width=True)
if st.button("Save Bounding Box"):
with open("temp_bbox.txt", "w") as bbox_file:
bbox_file.write(f"{xmin} {ymin} {xmax} {ymax}")
st.write(f"Bounding box saved to {os.path.abspath('temp_bbox.txt')}")
col1, col2 = st.columns(2)
with col2:
if st.button(
"Stop/Reset",
key="stop",
help="Stop the process and reset the app",
type="primary",
):
st.write("Stopping...")
os.system("rm -r ./temp_images")
os.system("rm ./temp_bbox.txt")
os.system("rm -r ./temp_processed_images")
os.system("rm ./temp_video.mp4")
os.system("rm ./temp_frame.jpg")
st.write("Process interrupted")
with col1:
if st.button(
"Segment Video", key="segment", help="Segment the video", type="secondary"
):
if use_bbox:
if not os.path.exists("./temp_bbox.txt"):
with open("temp_bbox.txt", "w") as bbox_file:
bbox_file.write(f"{xmin} {ymin} {xmax} {ymax}")
else:
with open("temp_bbox.txt", "w") as bbox_file:
bbox_file.write(f"0 0 {original_width} {original_height}")
st.write("Segmenting video...")
so = st.empty()
with rd.stdouterr(to=st.sidebar):
segment_video(
video_filename="temp_video.mp4",
dir_frames="temp_images",
image_start=0,
image_end=0,
bbox_file="temp_bbox.txt",
skip_vid2im=False,
mobile_sam_weights="./models/mobile_sam.pt",
auto_detect=not use_bbox,
background_color=background_color,
output_video="video_segmented.mp4",
output_dir="temp_processed_images",
pbar=False,
reverse_mask=not use_bbox,
)
os.system("rm -rf ./temp_images")
os.system("rm -rf ./temp_bbox.txt")
os.system("rm -rf ./temp_processed_images")
os.system("rm -rf ./temp_video.mp4")
st.video("./video_segmented.mp4")
st.write(f"Video saved to {os.path.abspath('video_segmented.mp4')}")
vid_file = open("video_segmented.mp4", "rb")
vid_bytes = vid_file.read()
st.download_button(
label="Download Segmented Video",
data=vid_bytes,
file_name="video_segmented.mp4",
)
vid_file.close()
|