Spaces:
Sleeping
Sleeping
| 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") | |