File size: 2,155 Bytes
292e401 1ab1e2e 292e401 |
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 |
import streamlit as st
import numpy as np
import cv2
from PIL import Image
# Load model files
prototxt_path = "colorization_deploy_v2.prototxt"
model_path = "colorization_release_v2.caffemodel"
kernel_path = "pts_in_hull.npy"
# Load the 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")]
# Streamlit App
st.title("Black-and-White Image Colorization")
st.write("Upload a black-and-white image to colorize it.")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Convert uploaded file to OpenCV format
image = Image.open(uploaded_file)
bw_image = np.array(image.convert("RGB"))
bw_image = cv2.cvtColor(bw_image, cv2.COLOR_RGB2BGR)
# Preprocessing for colorization
normalized = bw_image.astype("float32") / 255.0
lab = cv2.cvtColor(normalized, cv2.COLOR_BGR2LAB)
resized = cv2.resize(lab, (224, 224))
L = cv2.split(resized)[0]
L -= 50
# Predict color channels
net.setInput(cv2.dnn.blobFromImage(L))
ab = net.forward()[0, :, :, :].transpose((1, 2, 0))
ab = cv2.resize(ab, (bw_image.shape[1], bw_image.shape[0]))
# Combine L and ab channels
L = cv2.split(lab)[0]
colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)
# Convert LAB to BGR
colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2BGR)
colorized = (255 * colorized).astype("uint8")
# Display results
st.image(colorized, channels="BGR", caption="Colorized Image")
# Provide download link
colorized_image = Image.fromarray(cv2.cvtColor(colorized, cv2.COLOR_BGR2RGB))
colorized_image.save("colorized_output.jpg")
with open("colorized_output.jpg", "rb") as file:
btn = st.download_button(
label="Download colorized image",
data=file,
file_name="colorized_image.jpg",
mime="image/jpeg"
)
|