597 lines
22 KiB
Python
597 lines
22 KiB
Python
# this script will process multiple video segments (5 second clips)
|
|
# of a long (1 hour total video) and will greenscreen/key out everything
|
|
# except tracked objects.
|
|
#
|
|
# To specify which items are tracked, the user will call the script with
|
|
# --segments-collect-points 1 5 10 ... Which are segments to be treated
|
|
# as "keyframes", for these we will ask the user to select the point
|
|
# we want to track by hand, this is done for large changes in the objects
|
|
# position in the video.
|
|
#
|
|
# For other segments, since the object is mostly static in the frame,
|
|
# We will use the previous segments final frame (an image
|
|
# with the object we want to track visible, and everything else green)
|
|
# to determine an input mask and use add_new_mask() instead of selecting
|
|
# points.
|
|
#
|
|
# Each segment has 2 versions of each frame, on high quality used for
|
|
# final rendering and 1 low quality used to speed up inference
|
|
#
|
|
# When the script finishes, each segment should have an output directory
|
|
# with the same object tracked throughout the every frame in all the segment directories
|
|
#
|
|
# I will then turn these back into a video using ffmpeg but that is outside the scope
|
|
# of this program
|
|
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
import cupy as cp
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import torch
|
|
import logging
|
|
import sys
|
|
import gc
|
|
from sam2.build_sam import build_sam2_video_predictor
|
|
import argparse
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
|
|
|
|
# Variables for input and output directories
|
|
SAM2_CHECKPOINT = "../checkpoints/sam2.1_hiera_large.pt"
|
|
MODEL_CFG = "configs/sam2.1/sam2.1_hiera_l.yaml"
|
|
GREEN = [0, 255, 0]
|
|
BLUE = [255, 0, 0]
|
|
|
|
INFERENCE_SCALE = 0.25
|
|
FULL_SCALE = 1.0
|
|
|
|
def open_video(video_path):
|
|
"""
|
|
Opens a video file and returns a generator that yields frames.
|
|
|
|
Parameters:
|
|
- video_path: Path to the video file.
|
|
|
|
Returns:
|
|
- A generator that yields frames from the video.
|
|
"""
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
print(f"Error: Could not open video file {video_path}")
|
|
return
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
yield frame
|
|
cap.release()
|
|
|
|
def load_previous_segment_mask(prev_segment_dir):
|
|
mask_path = os.path.join(prev_segment_dir, "mask.png")
|
|
mask_image = cv2.imread(mask_path)
|
|
|
|
if mask_image is None:
|
|
raise FileNotFoundError(f"Mask image not found at {mask_path}")
|
|
|
|
# Ensure the mask_image has three color channels
|
|
if len(mask_image.shape) != 3 or mask_image.shape[2] != 3:
|
|
raise ValueError("Mask image does not have three color channels.")
|
|
|
|
mask_image = mask_image.astype(np.uint8)
|
|
|
|
# Extract Object A and Object B masks
|
|
mask_a = np.all(mask_image == GREEN, axis=2)
|
|
mask_b = np.all(mask_image == BLUE, axis=2)
|
|
|
|
per_obj_input_mask = {1: mask_a, 2: mask_b}
|
|
input_palette = None # No palette needed for binary mask
|
|
|
|
return per_obj_input_mask, input_palette
|
|
|
|
|
|
def apply_green_mask(frame, masks):
|
|
# Convert frame and masks to CuPy arrays
|
|
frame_gpu = cp.asarray(frame)
|
|
combined_mask = cp.zeros(frame_gpu.shape[:2], dtype=cp.bool_)
|
|
|
|
for mask in masks:
|
|
mask_gpu = cp.asarray(mask.squeeze())
|
|
if mask_gpu.shape != frame_gpu.shape[:2]:
|
|
resized_mask = cv2.resize(cp.asnumpy(mask_gpu).astype(cp.float32),
|
|
(frame_gpu.shape[1], frame_gpu.shape[0]))
|
|
mask_gpu = cp.asarray(resized_mask > 0.5) # Convert back to CuPy boolean array
|
|
else:
|
|
mask_gpu = mask_gpu.astype(cp.bool_) # Ensure boolean type
|
|
combined_mask |= mask_gpu # Perform the bitwise OR operation
|
|
|
|
green_background = cp.full(frame_gpu.shape, cp.array([0, 255, 0], dtype=cp.uint8), dtype=cp.uint8)
|
|
result_frame = cp.where(combined_mask[..., None], frame_gpu, green_background)
|
|
return cp.asnumpy(result_frame) # Convert back to NumPy
|
|
|
|
|
|
# def apply_green_mask(frame, masks):
|
|
# """
|
|
# Applies masks to the frame, replacing the background with green.
|
|
#
|
|
# Parameters:
|
|
#
|
|
# - frame: numpy array representing the image frame.
|
|
# - masks: list of numpy arrays representing the masks.
|
|
#
|
|
# Returns:
|
|
# - result_frame: numpy array with the green background applied.
|
|
# """
|
|
# # Initialize combined mask as a boolean array
|
|
# combined_mask = cp.zeros(frame.shape[:2], dtype=bool)
|
|
#
|
|
#
|
|
# for mask in masks:
|
|
# mask = mask.squeeze()
|
|
#
|
|
# # Resize the mask if necessary
|
|
# if mask.shape != frame.shape[:2]:
|
|
# # Resize the mask using bilinear interpolation
|
|
#
|
|
# # and convert it to float32 for accurate interpolation
|
|
# resized_mask = cv2.resize(
|
|
# mask.astype(cp.float32),
|
|
# (frame.shape[1], frame.shape[0]),
|
|
# interpolation=cv2.INTER_CUBIC
|
|
# )
|
|
# # Threshold the resized mask to obtain a boolean mask
|
|
# # add a small gausian blur to the mask to smooth out the edges
|
|
#
|
|
# mask = resized_mask > 0.5
|
|
# else:
|
|
# # Ensure mask is boolean
|
|
# mask = mask.astype(bool)
|
|
#
|
|
# # Combine masks using logical OR
|
|
# combined_mask |= mask # Now both arrays are bool
|
|
#
|
|
# # Create a green background image
|
|
# green_background = cp.full_like(frame, [0, 255, 0])
|
|
# # Use combined mask to overlay the original frame onto the green background
|
|
# result_frame = cp.where(
|
|
# combined_mask[..., None],
|
|
#
|
|
# frame,
|
|
# green_background
|
|
# )
|
|
# #result_frame = frame.copy()
|
|
# #result_frame[~combined_mask] = [0, 255, 0]
|
|
#
|
|
# return result_frame
|
|
|
|
def initialize_predictor():
|
|
if torch.cuda.is_available():
|
|
device = torch.device("cuda")
|
|
elif torch.backends.mps.is_available():
|
|
device = torch.device("mps")
|
|
print(
|
|
"\nSupport for MPS devices is preliminary. SAM 2 is trained with CUDA and might "
|
|
"give numerically different outputs and sometimes degraded performance on MPS."
|
|
)
|
|
# Enable MPS fallback for operations not supported on MPS
|
|
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
|
|
else:
|
|
device = torch.device("cpu")
|
|
logger.info(f"Using device: {device}")
|
|
predictor = build_sam2_video_predictor(MODEL_CFG, SAM2_CHECKPOINT, device=device)
|
|
return predictor
|
|
|
|
|
|
def load_first_frame(video_path, scale=1.0):
|
|
"""
|
|
Opens a video file and returns the first frame, scaled as specified.
|
|
|
|
Parameters:
|
|
- video_path: Path to the video file.
|
|
- scale: Scaling factor for the frame (default is 1.0 for original size).
|
|
|
|
Returns:
|
|
- first_frame: The first frame of the video, scaled accordingly.
|
|
"""
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
logger.error(f"Error: Could not open video file {video_path}")
|
|
return None
|
|
|
|
ret, frame = cap.read()
|
|
cap.release()
|
|
|
|
if not ret:
|
|
logger.error(f"Error: Could not read frame from video file {video_path}")
|
|
return None
|
|
|
|
if scale != 1.0:
|
|
frame = cv2.resize(
|
|
frame, None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR
|
|
)
|
|
|
|
return frame
|
|
|
|
def select_points_old(first_frame):
|
|
points_a = []
|
|
points_b = []
|
|
current_object = 'A'
|
|
point_count = 0
|
|
selection_complete = False
|
|
|
|
def mouse_callback(event, x, y, flags, param):
|
|
nonlocal points_a, points_b, current_object, point_count, selection_complete
|
|
if event == cv2.EVENT_LBUTTONDOWN:
|
|
if current_object == 'A':
|
|
points_a.append((x, y))
|
|
point_count += 1
|
|
print(f"Selected point {point_count} for Object A: ({x}, {y})")
|
|
if len(points_a) == 5: # Collect 4 points for Object A
|
|
current_object = 'B'
|
|
point_count = 0
|
|
print("Select point 1 for Object B")
|
|
elif current_object == 'B':
|
|
points_b.append((x, y))
|
|
point_count += 1
|
|
print(f"Selected point {point_count} for Object B: ({x}, {y})")
|
|
if len(points_b) == 5: # Collect 4 points for Object B
|
|
selection_complete = True
|
|
|
|
print("Select point 1 for Object A")
|
|
cv2.namedWindow('Select Points', cv2.WINDOW_NORMAL)
|
|
cv2.resizeWindow('Select Points', int(first_frame.shape[1] * (500 / first_frame.shape[0])), 500)
|
|
cv2.imshow('Select Points', first_frame)
|
|
cv2.setMouseCallback('Select Points', mouse_callback)
|
|
|
|
while not selection_complete:
|
|
cv2.waitKey(1)
|
|
|
|
cv2.destroyAllWindows()
|
|
return np.array(points_a, dtype=np.float32), np.array(points_b, dtype=np.float32)
|
|
|
|
def select_points(first_frame):
|
|
points_a = []
|
|
point_count = 0
|
|
selection_complete = False
|
|
frame_width = first_frame.shape[1]
|
|
half_frame_width = frame_width // 2 # Integer division for pixel coordinates
|
|
|
|
def mouse_callback(event, x, y, flags, param):
|
|
nonlocal points_a, point_count, selection_complete
|
|
if event == cv2.EVENT_LBUTTONDOWN:
|
|
points_a.append((x, y))
|
|
point_count += 1
|
|
print(f"Selected point {point_count} for Object A: ({x}, {y})")
|
|
if len(points_a) == 5: # Collect 5 points for Object A
|
|
selection_complete = True
|
|
|
|
print("Select 5 points for Object A (left side)")
|
|
cv2.namedWindow('Select Points', cv2.WINDOW_NORMAL)
|
|
cv2.resizeWindow('Select Points', int(first_frame.shape[1] * (500 / first_frame.shape[0])), 500)
|
|
cv2.imshow('Select Points', first_frame)
|
|
cv2.setMouseCallback('Select Points', mouse_callback)
|
|
|
|
while not selection_complete:
|
|
cv2.waitKey(1)
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
# Automatically generate points for Object B by shifting x-coordinates
|
|
points_a = np.array(points_a, dtype=np.float32)
|
|
points_b = points_a.copy()
|
|
points_b[:, 0] += half_frame_width # Shift x-coordinate by half the frame width
|
|
|
|
# Ensure that the shifted points are within the frame boundaries
|
|
points_b[:, 0] = np.clip(points_b[:, 0], 0, frame_width - 1)
|
|
|
|
return points_a, points_b
|
|
|
|
def add_points_to_predictor(predictor, inference_state, points, obj_id):
|
|
labels = np.array([1, 1, 1, 1, 1], np.int32) # Update labels to match 4 points
|
|
points = np.array(points, dtype=np.float32) # Ensure points have shape (4, 2)
|
|
try:
|
|
print(f"Adding points for Object {obj_id}: {points}")
|
|
_, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box(
|
|
inference_state=inference_state,
|
|
frame_idx=0,
|
|
obj_id=obj_id,
|
|
points=points,
|
|
labels=labels,
|
|
)
|
|
print(f"Object {obj_id} added successfully: {out_obj_ids}")
|
|
return out_mask_logits
|
|
except Exception as e:
|
|
print(f"Error adding points for Object {obj_id}: {e}")
|
|
exit()
|
|
|
|
def propagate_masks(predictor, inference_state):
|
|
video_segments = {}
|
|
for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
|
|
video_segments[out_frame_idx] = {
|
|
out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
|
|
for i, out_obj_id in enumerate(out_obj_ids)
|
|
}
|
|
return video_segments
|
|
|
|
def apply_colored_mask(frame, masks_a, masks_b):
|
|
colored_mask = np.zeros_like(frame)
|
|
|
|
# Apply colors to the masks
|
|
for mask in masks_a:
|
|
mask = mask.squeeze()
|
|
if mask.shape != frame.shape[:2]:
|
|
mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
indices = np.where(mask)
|
|
colored_mask[mask] = [0, 255, 0] # Green for Object A
|
|
|
|
for mask in masks_b:
|
|
mask = mask.squeeze()
|
|
if mask.shape != frame.shape[:2]:
|
|
mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
indices = np.where(mask)
|
|
colored_mask[mask] = [255, 0, 0] # Blue for Object B
|
|
|
|
return colored_mask
|
|
|
|
|
|
def process_and_save_output_video(video_path, output_video_path, video_segments, use_nvenc=False):
|
|
"""
|
|
Process high-resolution frames, apply upscaled masks, and save the output video.
|
|
"""
|
|
cap = cv2.VideoCapture(video_path)
|
|
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
fps = cap.get(cv2.CAP_PROP_FPS) or 59.94
|
|
|
|
# Setup VideoWriter with desired settings
|
|
if use_nvenc:
|
|
# Use FFmpeg with NVENC offloading for H.265 encoding
|
|
import subprocess
|
|
|
|
if sys.platform == 'darwin':
|
|
encoder = 'hevc_videotoolbox'
|
|
else:
|
|
encoder = 'hevc_nvenc'
|
|
|
|
command = [
|
|
'ffmpeg',
|
|
'-y', # Overwrite output file if it exists
|
|
'-f', 'rawvideo',
|
|
'-vcodec', 'rawvideo',
|
|
'-pix_fmt', 'bgr24',
|
|
'-s', f'{frame_width}x{frame_height}',
|
|
'-r', str(fps),
|
|
'-i', '-', # Input from stdin
|
|
'-an', # No audio
|
|
'-vcodec', encoder,
|
|
'-pix_fmt', 'nv12',
|
|
'-preset', 'slow',
|
|
'-b:v', '50M',
|
|
output_video_path
|
|
]
|
|
process = subprocess.Popen(command, stdin=subprocess.PIPE)
|
|
else:
|
|
# Use OpenCV VideoWriter
|
|
fourcc = cv2.VideoWriter_fourcc(*'HEVC') # H.265
|
|
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
|
|
|
|
frame_idx = 0
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret or frame_idx >= len(video_segments):
|
|
break
|
|
|
|
masks = [video_segments[frame_idx][out_obj_id] for out_obj_id in video_segments[frame_idx]]
|
|
upscaled_masks = []
|
|
|
|
for mask in masks:
|
|
mask = mask.squeeze()
|
|
upscaled_mask = cv2.resize(mask.astype(np.uint8), (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
upscaled_masks.append(upscaled_mask)
|
|
|
|
result_frame = apply_green_mask(frame, upscaled_masks)
|
|
|
|
# Write frame to output
|
|
if use_nvenc:
|
|
process.stdin.write(result_frame.tobytes())
|
|
else:
|
|
out.write(result_frame)
|
|
|
|
frame_idx += 1
|
|
|
|
cap.release()
|
|
if use_nvenc:
|
|
process.stdin.close()
|
|
process.wait()
|
|
else:
|
|
out.release()
|
|
|
|
def get_video_file_name(index):
|
|
return f"segment_{str(index).zfill(3)}.mp4"
|
|
|
|
def do_collect_segment_points(base_dir, segments, collect_points_segments, scale=1.0):
|
|
logger.info("Collecting points for requested segments.")
|
|
for i, segment in enumerate(segments):
|
|
segment_index = int(segment.split("_")[1])
|
|
segment_dir = os.path.join(base_dir, segment)
|
|
points_file = os.path.join(segment_dir, "segment_points")
|
|
video_file = os.path.join(segment_dir, get_video_file_name(i))
|
|
if segment_index in collect_points_segments and not os.path.exists(points_file):
|
|
first_frame = load_first_frame(video_file, scale)
|
|
points_a, points_b = select_points(first_frame)
|
|
with open(points_file, 'w') as f:
|
|
np.savetxt(f, points_a, header="Object A Points")
|
|
np.savetxt(f, points_b, header="Object B Points")
|
|
|
|
def save_final_masks(video_segments, mask_output_path):
|
|
"""
|
|
Save the final masks as a colored image.
|
|
"""
|
|
last_frame_idx = max(video_segments.keys())
|
|
masks_dict = video_segments[last_frame_idx]
|
|
# Assuming you have two objects with IDs 1 and 2
|
|
mask_a = masks_dict.get(1).squeeze()
|
|
mask_b = masks_dict.get(2).squeeze()
|
|
|
|
#create a black image with dimensions with shape (mask_a.y, mask_a.x, 3)
|
|
black_frame = np.zeros((mask_a.shape[0], mask_a.shape[1], 3), dtype=np.uint8)
|
|
if mask_a is None or mask_b is None:
|
|
print("Error: Masks for objects not found.")
|
|
return
|
|
|
|
#convert mask to np.uint8
|
|
mask_a = mask_a.astype(bool)
|
|
mask_b = mask_b.astype(bool)
|
|
|
|
# mask a
|
|
mask_a = mask_a.squeeze()
|
|
indices = np.where(mask_a)
|
|
black_frame[mask_a] = GREEN
|
|
# mask b
|
|
mask_b = mask_b.squeeze()
|
|
indices = np.where(mask_b)
|
|
black_frame[mask_b] = BLUE
|
|
|
|
# Save the mask image
|
|
cv2.imwrite(mask_output_path, black_frame)
|
|
|
|
def create_low_res_video(input_video_path, output_video_path, scale):
|
|
"""
|
|
Creates a low-resolution version of the input video for inference.
|
|
"""
|
|
cap = cv2.VideoCapture(input_video_path)
|
|
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) * scale)
|
|
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * scale)
|
|
fps = cap.get(cv2.CAP_PROP_FPS) or 59.94
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
low_res_frame = cv2.resize(frame, (frame_width, frame_height), interpolation=cv2.INTER_LINEAR)
|
|
out.write(low_res_frame)
|
|
|
|
cap.release()
|
|
out.release()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Process video segments.")
|
|
# arg for setting base_dir
|
|
parser.add_argument("--base-dir", type=str, help="Base directory for video segments.")
|
|
parser.add_argument("--segments-collect-points", nargs='+', type=int, help="Segments for which to collect points.")
|
|
args = parser.parse_args()
|
|
|
|
base_dir = args.base_dir
|
|
segments = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d)) and d.startswith("segment_")]
|
|
segments.sort(key=lambda x: int(x.split("_")[1]))
|
|
scaled_frames_dir_name = "frames_scaled"
|
|
fullres_frames_dir_name = "frames" # iwant to render the final video with these frames
|
|
|
|
collect_points_segments = args.segments_collect_points if args.segments_collect_points else []
|
|
#inference_scale for getting the mask, then use full scale when rendering the video
|
|
do_collect_segment_points(base_dir, segments, collect_points_segments, scale=INFERENCE_SCALE)
|
|
|
|
for i, segment in enumerate(segments):
|
|
segment_index = int(segment.split("_")[1])
|
|
segment_dir = os.path.join(base_dir, segment)
|
|
video_file_name = get_video_file_name(i)
|
|
video_path = os.path.join(segment_dir, video_file_name)
|
|
output_done_file = os.path.join(segment_dir, "output_frames_done")
|
|
if os.path.exists(output_done_file):
|
|
print(f"Segment {segment} already processed. Skipping.")
|
|
continue
|
|
|
|
logger.info(f"Processing segment {segment}")
|
|
|
|
# Initialize predictor
|
|
predictor = initialize_predictor()
|
|
|
|
# Prepare low-resolution video frames for inference
|
|
low_res_video_path = os.path.join(segment_dir, "low_res_video.mp4")
|
|
if not os.path.exists(low_res_video_path):
|
|
create_low_res_video(video_path, low_res_video_path, INFERENCE_SCALE)
|
|
logger.info(f"Low-resolution video created for segment {segment}")
|
|
else:
|
|
logger.info(f"Low-resolution video already exists for segment {segment}, reuse")
|
|
|
|
# Initialize inference state with low-resolution video
|
|
inference_state = predictor.init_state(video_path=low_res_video_path, async_loading_frames=True)
|
|
|
|
# Load points or previous masks
|
|
points_file = os.path.join(segment_dir, "segment_points")
|
|
if os.path.exists(points_file):
|
|
logger.info(f"Using segment_points for segment {segment}")
|
|
points = np.loadtxt(points_file, comments="#")
|
|
points_a = points[:5]
|
|
points_b = points[5:]
|
|
else:
|
|
points_a = points_b = None
|
|
|
|
collect_points = segment_index in collect_points_segments
|
|
|
|
if i > 0 and not collect_points and points_a is None:
|
|
# Try to load previous segment mask
|
|
logger.info(f"Using previous segment mask for segment {segment}")
|
|
prev_segment_dir = os.path.join(base_dir, segments[i - 1])
|
|
prev_mask_path = os.path.join(prev_segment_dir, "mask.png")
|
|
if os.path.exists(prev_mask_path):
|
|
per_obj_input_mask, input_palette = load_previous_segment_mask(prev_segment_dir)
|
|
# Add previous masks to predictor
|
|
for obj_id, mask in per_obj_input_mask.items():
|
|
predictor.add_new_mask(inference_state, 0, obj_id, mask)
|
|
else:
|
|
print(f"Warning: Previous segment mask not found for segment {segment}.")
|
|
continue # Skip this segment or handle as needed
|
|
else:
|
|
if points_a is not None and points_b is not None:
|
|
print("Using points for segment")
|
|
# Add points to predictor
|
|
add_points_to_predictor(predictor, inference_state, points_a, obj_id=1)
|
|
add_points_to_predictor(predictor, inference_state, points_b, obj_id=2)
|
|
else:
|
|
print("Error: No points available for segment.")
|
|
continue # Skip this segment
|
|
|
|
# Perform inference and collect masks per frame
|
|
video_segments = propagate_masks(predictor, inference_state)
|
|
|
|
# Process high-resolution frames and save output video
|
|
output_video_path = os.path.join(segment_dir, f"output_{segment_index}.mp4")
|
|
print("Processing of segment complete, attempting to save process full video from low res masks")
|
|
process_and_save_output_video(
|
|
video_path,
|
|
output_video_path,
|
|
video_segments,
|
|
use_nvenc=True # Set to True to use NVENC offloading
|
|
)
|
|
|
|
# Save final masks
|
|
mask_output_path = os.path.join(segment_dir, "mask.png")
|
|
save_final_masks(video_segments, mask_output_path)
|
|
|
|
# Clean up
|
|
predictor.reset_state(inference_state)
|
|
del inference_state
|
|
del video_segments
|
|
del predictor
|
|
gc.collect()
|
|
|
|
try:
|
|
os.remove(low_res_video_path)
|
|
logger.info(f"Deleted low-resolution video for segment {segment}")
|
|
except Exception as e:
|
|
logger.warning(f"Could not delete low-resolution video for segment {segment}: {e}")
|
|
|
|
open(output_done_file, 'a').close()
|
|
print("Processing complete.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|