Spaces:
Sleeping
Sleeping
File size: 3,635 Bytes
ff3e89b | 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 | import cv2
import numpy as np
import torch
import streamlit as st
from PIL import Image
import colorizers
# Load pretrained colorization model
model = colorizers.siggraph17(pretrained=True).eval()
# Session state init
if 'processed_image' not in st.session_state:
st.session_state.processed_image = None
if 'original_image' not in st.session_state:
st.session_state.original_image = None
if 'history' not in st.session_state:
st.session_state.history = []
# Convert OpenCV image to PIL
def display_image_cv2(image):
rgb_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
return Image.fromarray(rgb_img)
# Colorization logic
def colouring_image(file, model):
img = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_GRAYSCALE)
original = cv2.cvtColor(cv2.resize(img, (256, 256)), cv2.COLOR_GRAY2BGR)
img = cv2.resize(img, (256, 256)) / 255.0 * 100
img_tensor = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float()
with torch.no_grad():
ab = model(img_tensor).cpu().numpy()[0].transpose((1, 2, 0))
lab = np.concatenate((img[:, :, np.newaxis], ab), axis=2)
bgr = cv2.cvtColor(lab.astype(np.float32), cv2.COLOR_Lab2BGR)
bgr = np.clip(bgr * 255, 0, 255).astype(np.uint8)
return bgr, original
# UI Setup
st.set_page_config(page_title="Image Colorizer", layout="wide")
st.title("🎨 Image Colorization and Post-Processing Tool")
uploaded_file = st.file_uploader("Upload a grayscale image", type=["jpg", "jpeg", "png", "bmp"])
if uploaded_file:
colorized, original = colouring_image(uploaded_file, model)
st.session_state.processed_image = colorized.copy()
st.session_state.original_image = original
st.session_state.history = [colorized.copy()]
st.subheader("Preview:")
col1, col2 = st.columns(2)
with col1:
st.image(display_image_cv2(original), caption="Original Image", use_container_width=True)
with col2:
st.image(display_image_cv2(colorized), caption="Colorized Image", use_container_width=True)
st.markdown("---")
# Button row
colA, colB, colC, colD = st.columns(4)
with colA:
if st.button("🔪 Sharpen"):
kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
sharpened = cv2.filter2D(st.session_state.processed_image, -1, kernel)
st.session_state.history.append(st.session_state.processed_image.copy())
st.session_state.processed_image = sharpened
st.image(display_image_cv2(sharpened), caption="Sharpened Image", use_container_width=True)
with colB:
if st.button("💧 Blur"):
blurred = cv2.GaussianBlur(st.session_state.processed_image, (15, 15), 0)
st.session_state.history.append(st.session_state.processed_image.copy())
st.session_state.processed_image = blurred
st.image(display_image_cv2(blurred), caption="Blurred Image", use_container_width=True)
with colC:
if st.button("↩️ Undo"):
if len(st.session_state.history) > 1:
st.session_state.history.pop()
st.session_state.processed_image = st.session_state.history[-1]
st.image(display_image_cv2(st.session_state.processed_image), caption="Undo Applied", use_container_width=True)
else:
st.warning("Nothing to undo.")
with colD:
if st.session_state.processed_image is not None:
buf = cv2.imencode(".jpg", st.session_state.processed_image)[1].tobytes()
st.download_button(label="💾 Save Image", data=buf, file_name="colorized.jpg", mime="image/jpeg")
|