tarinmodel3 / app.py
nagasurendra's picture
Update app.py
54d292a verified
raw
history blame
1.64 kB
import torch
import cv2
import numpy as np
# Load your model (assuming it is a PyTorch model)
model = torch.load('.data/model.pt')
model.eval()
# Open video file (input video)
input_video = cv2.VideoCapture('input_video.mp4')
# Get the frame width, height, and frames per second (fps) from the input video
frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = input_video.get(cv2.CAP_PROP_FPS)
# Define the output video writer
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # You can change this to any codec
output_video = cv2.VideoWriter('output_video.mp4', fourcc, fps, (frame_width, frame_height))
while True:
# Read a frame from the input video
ret, frame = input_video.read()
if not ret:
break # End of video
# Preprocess the frame if necessary (depends on your model)
# For example, convert to tensor and normalize if required
frame_tensor = torch.tensor(frame).float().unsqueeze(0) # Add batch dimension
# Pass the frame through the model
with torch.no_grad():
output = model(frame_tensor) # Adjust based on your model's requirements
# Postprocess the output if necessary (depends on your model's output format)
output_frame = output.squeeze(0).cpu().numpy() # Remove batch dimension and convert to NumPy
# Convert the model output to a valid image format (if necessary)
output_frame = np.uint8(output_frame)
# Write the frame to the output video
output_video.write(output_frame)
# Release resources
input_video.release()
output_video.release()
cv2.destroyAllWindows()