File size: 3,280 Bytes
bc57072 7a6ffb4 bc57072 7a6ffb4 bc57072 7a6ffb4 ffea17c bc57072 ffea17c 7a6ffb4 bc57072 ff1d39f bc57072 ffea17c 7a6ffb4 bc57072 2a3d0b1 bc57072 2a3d0b1 bc57072 7a6ffb4 |
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 |
import streamlit as st
import numpy as np
import cv2
import tempfile
import os
# Load model files
prototxt_path = "colorization_deploy_v2.prototxt"
model_path = "colorization_release_v2.caffemodel"
kernel_path = "pts_in_hull.npy"
# Streamlit app title
st.title("Video Colorization App")
# File upload
uploaded_video = st.file_uploader("Upload a black and white video", type=["mp4", "avi"])
if uploaded_video is not None:
# Save uploaded video to a temporary file
tfile = tempfile.NamedTemporaryFile(delete=False)
tfile.write(uploaded_video.read())
video_path = tfile.name
# Output path for the colorized video
output_path = os.path.join(tempfile.gettempdir(), "colorized_video.mp4")
# Load the pre-trained model
net = cv2.dnn.readNetFromCaffe(prototxt_path, model_path)
points = np.load(kernel_path)
points = points.transpose().reshape(2, 313, 1, 1)
net.getLayer(net.getLayerId("class8_ab")).blobs = [points.astype(np.float32)]
net.getLayer(net.getLayerId("conv8_313_rh")).blobs = [np.full([1, 313], 2.686, dtype="float32")]
# Open the video file
cap = cv2.VideoCapture(video_path)
# Get video properties
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)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# Create a VideoWriter object to save the colorized video
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
# Initialize progress bar and frame counter
frame_count = 0
progress_bar = st.progress(0)
progress_text = st.empty() # Placeholder for frame count text
# Process each frame
while True:
ret, frame = cap.read()
if not ret:
break
frame_count += 1
progress_text.text(f"Processing frame {frame_count} of {total_frames}")
# Convert frame to LAB color space and preprocess
normalized = frame.astype("float32") / 255.0
lab = cv2.cvtColor(normalized, cv2.COLOR_BGR2LAB)
resized = cv2.resize(lab, (224, 224))
L = cv2.split(resized)[0]
L -= 43
# Set the input and get the colorization
net.setInput(cv2.dnn.blobFromImage(L))
ab = net.forward()[0, :, :, :].transpose((1, 2, 0))
ab = cv2.resize(ab, (frame.shape[1], frame.shape[0]))
# Combine with the L channel
L = cv2.split(lab)[0]
colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)
colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2BGR)
colorized = (255 * colorized).astype("uint8")
# Update the progress line and frame count display
progress_bar.progress(frame_count / total_frames)
# Write colorized frame to output
out.write(colorized)
# Release resources
cap.release()
out.release()
# Display the colorized video
st.success("Video colorization completed!")
st.video(output_path)
# Provide a download link for the colorized video
with open(output_path, "rb") as file:
st.download_button(label="Download Colorized Video", data=file, file_name="colorized_video.mp4", mime="video/mp4")
|