import os
from turtle import back
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
import numpy as np
import cv2
from glob import glob
from tqdm import tqdm
import tensorflow as tf
from tensorflow.keras.utils import CustomObjectScope
from metrics import dice_loss, dice_coef, iou
import moviepy.editor
from moviepy.editor import *
""" Global parameters """
H = 512
W = 512
""" Creating a directory """
def create_dir(path):
if not os.path.exists(path):
os.makedirs(path)
if __name__ == "__main__":
""" Seeding """
np.random.seed(42)
tf.random.set_seed(42)
""" Directory for storing files """
create_dir("processed_videos")
create_dir("frames")
""" Loading model: DeepLabV3+ """
with CustomObjectScope({'iou': iou, 'dice_coef': dice_coef, 'dice_loss': dice_loss}):
model = tf.keras.models.load_model("model.h5")
# model.summary()
""" Video Path """
video_path = "videos/in1.mp4"
# """Extracting the audio from video"""
# # Replace the parameter with the location of the video
# video = moviepy.editor.VideoFileClip(video_path)
# audio = video.audio
# # Replace the parameter with the location along with filename
# audio.write_audiofile("/Users/sotsys207/Downloads/Remove-Background-from-Video-using-TensorFlow-main/audio/sample.mp3")
""" Reading frames """
vs = cv2.VideoCapture(video_path)
_, frame = vs.read()
h, w, _ = frame.shape
vs.release()
fourcc = cv2.VideoWriter_fourcc('M','J','P','G')
out = cv2.VideoWriter(f'processed_videos/output.avi', fourcc, 30, (w, h), True)
cap = cv2.VideoCapture(video_path)
idx = 0
while True:
ret, frame = cap.read()
if ret == False:
cap.release()
out.release()
break
h, w, _ = frame.shape
ori_frame = frame
frame = cv2.resize(frame, (W, H))
frame = np.expand_dims(frame, axis=0)
frame = frame / 255.0
mask = model.predict(frame)[0]
mask = cv2.resize(mask, (w, h))
mask = mask > 0.5
mask = mask.astype(np.float32)
mask = np.expand_dims(mask, axis=-1)
photo_mask = mask
background_mask = np.abs(1-mask)
masked_frame = ori_frame * photo_mask
print(masked_frame.shape)
masked_frame = np.expand_dims(masked_frame, axis=0)
background_mask = np.concatenate([background_mask, background_mask, background_mask, background_mask], axis=-3)
# print(background_mask)
background_mask = background_mask * [255, 255, 255, 0]
print(background_mask.shape)
final_frame = masked_frame + background_mask
final_frame = final_frame.astype(np.uint8)
# print(final_frame)
cv2.imwrite(f"frames/{idx}.png", final_frame)
idx += 1
out.write(final_frame)
I am trying to add transparent background but i don't know how to do it can any one help me to do it.
I have tried it but it gives me error
1/1 [==============================] - 3s 3s/step
(1280, 720, 3)
(5120, 720, 4)
Traceback (most recent call last):
File "/Users/sotsys207/Downloads/Remove-Background-from-Video-using-TensorFlow-main/run.py", line 90, in <module>
final_frame = masked_frame + background_mask
ValueError: operands could not be broadcast together with shapes (1,1280,720,3) (5120,720,4)