S2IA / src /img2skt.py
Zeyu Zhao
Changes.
505e263
import os
import cv2
import gradio as gr
import numpy as np
from PIL import Image, ImageSequence
def image_to_sketch_gif(input_image: Image.Image):
# Convert PIL image to OpenCV format
open_cv_image = np.array(input_image.convert("RGB"))
open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)
# Convert to grayscale
grayscale_image = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2GRAY)
# Apply stronger Gaussian blur to smooth fine details
blurred_image = cv2.GaussianBlur(grayscale_image, (13, 13), 0)
# Use Canny Edge Detection with higher thresholds to detect only strong edges
edges = cv2.Canny(blurred_image, threshold1=100, threshold2=200)
# Dilate edges to make lines thicker
kernel = np.ones((5, 5), np.uint8)
edges = cv2.dilate(edges, kernel, iterations=1)
# Ensure binary format
_, binary_sketch = cv2.threshold(edges, 128, 255, cv2.THRESH_BINARY)
# Find connected components
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
binary_sketch, connectivity=8
)
# Sort components by size (excluding the background) and ignore very small ones
min_area = 250 # Minimum area to keep a component
components = sorted(
[
(i, stats[i, cv2.CC_STAT_AREA])
for i in range(1, num_labels)
if stats[i, cv2.CC_STAT_AREA] > min_area
],
key=lambda x: x[1],
reverse=True,
)
# Initialize an empty canvas for accumulation
accumulated_image = np.zeros_like(binary_sketch, dtype=np.uint8)
# Store frames
frames = []
for label, _ in components:
# Add the current component to the accumulation
accumulated_image[labels == label] = 255
# Convert OpenCV image to PIL image and append to frames
pil_frame = Image.fromarray(255 - accumulated_image)
frames.append(pil_frame.copy())
if not frames:
# Handle edge case where no components remain
frames.append(Image.fromarray(255 - accumulated_image))
return (
frames[0],
frames,
gr.Slider(
value=len(frames) - 1,
maximum=len(frames) - 1,
),
gr.Button(interactive=False),
)