|
import gradio as gr |
|
import cv2 |
|
import numpy as np |
|
import svgwrite |
|
from sklearn.cluster import KMeans |
|
|
|
|
|
def load_image(image_path): |
|
img = cv2.imread(image_path) |
|
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
|
|
|
|
def segment_by_color(image, n_clusters=5): |
|
|
|
pixels = image.reshape((-1, 3)) |
|
|
|
kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(pixels) |
|
|
|
segmented_img = kmeans.cluster_centers_[kmeans.labels_].reshape(image.shape).astype(np.uint8) |
|
return segmented_img, kmeans.labels_.reshape(image.shape[:2]) |
|
|
|
|
|
def generate_svg(image, labels, filename="output.svg"): |
|
dwg = svgwrite.Drawing(filename, profile='tiny') |
|
height, width, _ = image.shape |
|
for y in range(height): |
|
for x in range(width): |
|
color = image[y, x] |
|
label = labels[y, x] |
|
dwg.add(dwg.rect(insert=(x, y), size=(1, 1), fill=svgwrite.rgb(*color), stroke=svgwrite.rgb(*color))) |
|
dwg.save() |
|
|
|
|
|
def process_image(image, n_clusters): |
|
img = load_image(image.name) |
|
segmented, labels = segment_by_color(img, n_clusters) |
|
generate_svg(segmented, labels) |
|
return "output.svg" |
|
|
|
|
|
iface = gr.Interface( |
|
fn=process_image, |
|
inputs=[ |
|
gr.Image(type="file", label="アップロード画像"), |
|
gr.Slider(minimum=2, maximum=10, step=1, default=5, label="クラスタ数") |
|
], |
|
outputs="file", |
|
live=True |
|
) |
|
|
|
iface.launch() |
|
|