Update app.py

app.py

@@ -5,190 +5,132 @@ import numpy as np
 import cv2
 import vtracer
 
-def kmeans_segment(arr_rgba, n_colors, shadow_strength):
-    """Segmenta imagem em regioes de cor uniforme via K-means."""
-    r,g,b,a = arr_rgba[:,:,0],arr_rgba[:,:,1],arr_rgba[:,:,2],arr_rgba[:,:,3]
+def get_dominant_colors(arr_rgba, n_colors, shadow_strength):
+    r,g,b,a = arr_rgba[:,:,0].astype(np.float32),arr_rgba[:,:,1].astype(np.float32),arr_rgba[:,:,2].astype(np.float32),arr_rgba[:,:,3].astype(np.float32)
     opaque = a > 128
-
-    # Remove sombras
-    cmax = np.maximum(np.maximum(r,g),b).astype(np.float32)
-    cmin = np.minimum(np.minimum(r,g),b).astype(np.float32)
+    cmax = np.maximum(np.maximum(r,g),b)
+    cmin = np.minimum(np.minimum(r,g),b)
     sat = np.where(cmax>0,(cmax-cmin)/cmax,0)
-    lum = (0.299*r+0.587*g+0.114*b).astype(np.float32)/255
-    shadow = (sat < shadow_strength*0.35) & opaque
-
-    arr = arr_rgba.copy().astype(np.float32)
-    for ch in range(3):
-        arr[:,:,ch] = np.where(shadow & opaque, np.where(lum>0.5,255,0), arr[:,:,ch])
-    arr[:,:,3] = np.where(a>128,255,0)
-
-    # K-means nos pixels opacos e nao-sombra
-    valid = (arr[:,:,3]>128) & ~shadow
-    pixels = arr[:,:,:3][valid].astype(np.float32)
+    lum = (0.299*r+0.587*g+0.114*b)/255
+    shadow = sat < (shadow_strength*0.35)
+    valid = opaque & ~shadow
+    pixels = arr_rgba[:,:,:3][valid].astype(np.float32)
     if len(pixels) < n_colors:
-        return arr.astype(np.uint8), []
-
-    sample = pixels[::max(1,len(pixels)//10000)]
-    criteria = (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 30, 1.0)
-    k = min(n_colors, len(sample))
-    _, labels_s, centers = cv2.kmeans(sample, k, None, criteria, 5, cv2.KMEANS_PP_CENTERS)
-
-    # Aplica K-means a TODOS os pixels opacos
-    flat = arr[:,:,:3].reshape(-1,3).astype(np.float32)
-    flat_opaque = arr[:,:,3].reshape(-1) > 128
-
-    # Distancia para cada centro
-    dists = np.sqrt(((flat[:,np.newaxis,:]-centers[np.newaxis,:,:])**2).sum(axis=2))
-    best = np.argmin(dists,axis=1)
-
-    # Substitui cada pixel pela cor do centro mais proximo
-    result = arr.copy()
-    h,w = arr.shape[:2]
-    mapped = centers[best].reshape(h,w,3)
-    result[:,:,0] = np.where(arr[:,:,3]>128, mapped[:,:,0], arr[:,:,0])
-    result[:,:,1] = np.where(arr[:,:,3]>128, mapped[:,:,1], arr[:,:,1])
-    result[:,:,2] = np.where(arr[:,:,3]>128, mapped[:,:,2], arr[:,:,2])
-
-    # Ordena centros por frequencia
-    counts = np.bincount(best[flat_opaque])
+        return []
+    criteria = (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER,30,1.0)
+    k = min(n_colors, len(pixels))
+    sample = pixels[::max(1,len(pixels)//8000)]
+    _,labels,centers = cv2.kmeans(sample,k,None,criteria,5,cv2.KMEANS_PP_CENTERS)
+    counts = np.bincount(labels.flatten())
     order = np.argsort(-counts)
-    dominant = [centers[i] for i in order if 0.299*centers[i][0]/255+0.587*centers[i][1]/255+0.114*centers[i][2]/255 > 0.05]
-
-    return result.astype(np.uint8), dominant, centers[best].reshape(h,w,3), best.reshape(h,w)
-
-
-def clean_region(mask, close_iter=3, open_iter=1):
-    """Fecha buracos e remove ruido numa mascara binaria."""
-    k5 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))
-    k3 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3))
-    m = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, k5, iterations=close_iter)
-    m = cv2.morphologyEx(m,    cv2.MORPH_OPEN,  k3, iterations=open_iter)
-    return m
-
-
-def contours_to_svg_path(contour, epsilon_factor=0.002):
-    """Converte contorno OpenCV em path SVG suave."""
-    # Aproxima o contorno com menos pontos
-    epsilon = epsilon_factor * cv2.arcLength(contour, True)
-    approx = cv2.approxPolyDP(contour, epsilon, True)
-    if len(approx) < 3:
-        return None
-    pts = approx.reshape(-1,2)
-    d = f"M {pts[0][0]:.1f} {pts[0][1]:.1f}"
-    for pt in pts[1:]:
-        d += f" L {pt[0]:.1f} {pt[1]:.1f}"
-    d += " Z"
-    return d
-
-
-def rgb_to_hex(c):
-    return f"#{int(c[0]):02x}{int(c[1]):02x}{int(c[2]):02x}"
-
-
-def vectorize_with_contours(arr_rgba, dominant, label_map, centers, min_area, epsilon):
-    """Vetoriza cada regiao de cor usando contornos OpenCV."""
-    h, w = arr_rgba.shape[:2]
-    alpha = arr_rgba[:,:,3]
-    svg_paths = []
-
-    for ci, color in enumerate(dominant):
-        # Mascara desta cor (onde label_map == ci)
-        # Acha qual indice de centro corresponde a esta cor
-        dists = np.sqrt(((centers - color)**2).sum(axis=1))
-        center_idx = np.argmin(dists)
-
-        mask = ((label_map == center_idx) & (alpha > 128)).astype(np.uint8) * 255
-
-        if mask.sum() < min_area * 255:
-            continue
-
-        # Fecha buracos
-        mask_clean = clean_region(mask)
-
-        # Encontra contornos
-        contours, _ = cv2.findContours(mask_clean, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
-
-        color_hex = rgb_to_hex(color)
-        for cnt in contours:
-            area = cv2.contourArea(cnt)
-            if area < min_area:
-                continue
-            path_d = contours_to_svg_path(cnt, epsilon)
-            if path_d:
-                svg_paths.append(f'<path fill="{color_hex}" d="{path_d}"/>')
-
-    return svg_paths
-
-
-def preprocess(pil_img, shadow_strength, n_colors):
+    result = []
+    for i in order:
+        c = centers[i]
+        l = 0.299*c[0]/255+0.587*c[1]/255+0.114*c[2]/255
+        if l > 0.08:
+            result.append(c)
+    return result
+
+def preprocess_clean(pil_img, shadow_strength, n_colors):
     img = pil_img.convert('RGBA')
-    w, h = img.size
-    MAX, MIN = 1200, 500
+    w,h = img.size
+    MAX,MIN = 1200,500
     s = 1.0
-    if max(w,h) > MAX: s = MAX/max(w,h)
-    elif max(w,h) < MIN: s = MIN/max(w,h)
-    if s != 1.0:
-        img = img.resize((int(w*s), int(h*s)), Image.LANCZOS)
-
-    arr = np.array(img)
-    result, dominant, mapped, label_map = kmeans_segment(arr, int(n_colors), shadow_strength)
-    return Image.fromarray(result,'RGBA'), result, dominant, mapped, label_map
+    if max(w,h)>MAX: s=MAX/max(w,h)
+    elif max(w,h)<MIN: s=MIN/max(w,h)
+    if s!=1.0:
+        img = img.resize((int(w*s),int(h*s)),Image.LANCZOS)
+
+    arr = np.array(img,dtype=np.float32)
+    r,g,b,a = arr[:,:,0],arr[:,:,1],arr[:,:,2],arr[:,:,3]
+    opaque = a>128
+
+    if shadow_strength>0:
+        cmax = np.maximum(np.maximum(r,g),b)
+        cmin = np.minimum(np.minimum(r,g),b)
+        sat = np.where(cmax>0,(cmax-cmin)/cmax,0)
+        lum = (0.299*r+0.587*g+0.114*b)/255
+        shadow = (sat<shadow_strength*0.35)&opaque
+        for ch in range(3):
+            arr[:,:,ch] = np.where(shadow,np.where(lum>0.5,255,0),arr[:,:,ch])
 
+    arr[:,:,3] = np.where(a>128,255,0)
+    opaque = arr[:,:,3]>128
+
+    dominant = get_dominant_colors(arr.astype(np.uint8),int(n_colors),shadow_strength)
+    if dominant:
+        dom_arr = np.array(dominant,dtype=np.float32)
+        h2,w2 = arr.shape[:2]
+        flat = arr[:,:,:3].reshape(-1,3).astype(np.float32)
+        dists = np.sqrt(((flat[:,np.newaxis,:]-dom_arr[np.newaxis,:,:])**2).sum(axis=2))
+        best_idx = np.argmin(dists,axis=1)
+        best_dist = dists[np.arange(len(flat)),best_idx]
+        best_color = dom_arr[best_idx]
+        snap = best_dist<25
+        flat_result = flat.copy()
+        flat_result[snap] = best_color[snap]
+        arr[:,:,:3] = flat_result.reshape(h2,w2,3)
+
+    alpha = arr[:,:,3].astype(np.uint8)
+    k3 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3))
+    k5 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))
+    alpha = cv2.morphologyEx(alpha,cv2.MORPH_CLOSE,k5,iterations=3)
+    alpha = cv2.morphologyEx(alpha,cv2.MORPH_OPEN,k3,iterations=1)
+    alpha = np.where(alpha>127,255,0).astype(np.uint8)
+    arr[:,:,3] = alpha
+
+    if dominant:
+        dom_arr = np.array(dominant,dtype=np.float32)
+        rgb = arr[:,:,:3].astype(np.float32)
+        for color in dominant:
+            dist = np.sqrt(((rgb-color)**2).sum(axis=2))
+            mask = (dist<20).astype(np.uint8)*255
+            mask_closed = cv2.morphologyEx(mask,cv2.MORPH_CLOSE,k5,iterations=3)
+            new_pixels = (mask_closed>127)&(mask<127)&(alpha>127)
+            arr[:,:,0] = np.where(new_pixels,color[0],arr[:,:,0])
+            arr[:,:,1] = np.where(new_pixels,color[1],arr[:,:,1])
+            arr[:,:,2] = np.where(new_pixels,color[2],arr[:,:,2])
+
+    return Image.fromarray(arr.astype(np.uint8),'RGBA')
 
 def vectorize_image(
     input_image, shadow_strength, n_colors,
-    use_contours, min_area, epsilon,
     filter_speckle, color_precision, corner_threshold,
     length_threshold, path_precision,
 ):
     try:
         if input_image is None:
             raise gr.Error("Nenhuma imagem enviada")
-
         pil_img = input_image.convert('RGBA')
-        w, h = pil_img.size
+        w,h = pil_img.size
         print(f"Vetorizando {w}x{h}...")
         if w>2000 or h>2000:
             raise gr.Error("Imagem muito grande. Maximo: 2000x2000 px.")
 
-        pil_clean, arr_clean, dominant, mapped, label_map = preprocess(pil_img, shadow_strength, int(n_colors))
-        pw, ph = pil_clean.size
-        print(f"Pre-processamento: {pw}x{ph}, {len(dominant)} cores")
-
-        centers_arr = np.array([d for d in dominant], dtype=np.float32)
-
-        if use_contours and len(dominant) > 0:
-            # Modo contornos: paths limpos ponto a ponto
-            svg_paths = vectorize_with_contours(
-                arr_clean, dominant, label_map, centers_arr,
-                min_area=int(min_area), epsilon=float(epsilon)/1000
-            )
-
-            if svg_paths:
-                svg_str = f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {pw} {ph}" width="100%" height="100%">\n'
-                svg_str += '\n'.join(svg_paths)
-                svg_str += '\n</svg>'
-            else:
-                raise gr.Error("Nenhum contorno encontrado. Tente ajustar os parametros.")
-        else:
-            # Fallback: VTracer na imagem pre-processada
-            buf = io.BytesIO()
-            pil_clean.save(buf, format='PNG')
-            svg_str = vtracer.convert_raw_image_to_svg(
-                buf.getvalue(), img_format='png',
-                colormode='color', hierarchical='stacked', mode='spline',
-                filter_speckle=int(filter_speckle),
-                color_precision=int(color_precision),
-                layer_difference=16,
-                corner_threshold=int(corner_threshold),
-                length_threshold=float(length_threshold),
-                max_iterations=10, splice_threshold=60,
-                path_precision=int(path_precision),
-            )
-            if 'viewBox' not in svg_str:
-                svg_str = svg_str.replace('<svg ', f'<svg viewBox="0 0 {pw} {ph}" ', 1)
-            svg_str = re.sub(r'(<svg[^>]*?)width="[^"]*"', r'\1width="100%"', svg_str)
-            svg_str = re.sub(r'(<svg[^>]*?)height="[^"]*"', r'\1height="100%"', svg_str)
+        pil_img = preprocess_clean(pil_img,shadow_strength,n_colors)
+        pw,ph = pil_img.size
+        print(f"Apos pre-processamento: {pw}x{ph}")
+
+        buf = io.BytesIO()
+        pil_img.save(buf,format='PNG')
+        img_bytes = buf.getvalue()
+
+        svg_str = vtracer.convert_raw_image_to_svg(
+            img_bytes, img_format='png',
+            colormode='color', hierarchical='stacked', mode='spline',
+            filter_speckle=int(filter_speckle),
+            color_precision=int(color_precision),
+            layer_difference=16,
+            corner_threshold=int(corner_threshold),
+            length_threshold=float(length_threshold),
+            max_iterations=10, splice_threshold=60,
+            path_precision=int(path_precision),
+        )
+
+        if 'viewBox' not in svg_str:
+            svg_str = svg_str.replace('<svg ',f'<svg viewBox="0 0 {pw} {ph}" ',1)
+        svg_str = re.sub(r'(<svg[^>]*?)width="[^"]*"',r'\1width="100%"',svg_str)
+        svg_str = re.sub(r'(<svg[^>]*?)height="[^"]*"',r'\1height="100%"',svg_str)
 
         svg_path = f"/tmp/result_{os.getpid()}.svg"
         with open(svg_path,'w',encoding='utf-8') as f:
@@ -197,7 +139,7 @@ def vectorize_image(
         png_path = None
         try:
             import cairosvg
-            png_bytes = cairosvg.svg2png(bytestring=svg_str.encode(), scale=1.0)
+            png_bytes = cairosvg.svg2png(bytestring=svg_str.encode(),scale=1.0)
             png_path = f"/tmp/result_{os.getpid()}.png"
             with open(png_path,'wb') as f: f.write(png_bytes)
         except Exception as e:
@@ -211,44 +153,32 @@ def vectorize_image(
     except Exception as e:
         raise gr.Error(str(e))
 
-
 with gr.Blocks(title="Iluminados Vectorizer") as demo:
-    gr.Markdown("## Iluminados Vectorizer - OpenCV + VTracer")
-
+    gr.Markdown("## Iluminados Vectorizer - VTracer API")
     with gr.Row():
         with gr.Column():
             inp = gr.Image(label="Imagem de entrada", type="pil")
-
             with gr.Accordion("Pre-processamento", open=True):
-                shadow_strength = gr.Slider(0,1.0, value=0.6, step=0.05, label="Remover sombras")
-                n_colors        = gr.Slider(2,32,  value=16,  step=1,    label="Numero de cores")
-
-            with gr.Accordion("Modo de vetorizacao", open=True):
-                use_contours = gr.Checkbox(value=True, label="Usar contornos OpenCV (linhas mais limpas)")
-                min_area     = gr.Slider(10,2000, value=100, step=10,  label="Area minima de regiao")
-                epsilon      = gr.Slider(1,20,    value=4,   step=0.5, label="Suavidade dos contornos")
-
-            with gr.Accordion("VTracer (fallback)", open=False):
-                filter_speckle   = gr.Slider(1,100,value=6,  step=1,   label="Filtro fragmentos")
+                shadow_strength = gr.Slider(0,1.0,value=0.65,step=0.05,label="Remover sombras e gradientes")
+                n_colors        = gr.Slider(4,32, value=16,  step=1,   label="Numero de cores")
+            with gr.Accordion("Vetorizacao", open=False):
+                filter_speckle   = gr.Slider(1,100,value=6,  step=1,   label="Filtro de fragmentos")
                 color_precision  = gr.Slider(1,8,  value=6,  step=1,   label="Precisao de cor")
                 corner_threshold = gr.Slider(1,180,value=60, step=1,   label="Limiar de canto")
-                length_threshold = gr.Slider(0.5,10.,value=3.5,step=0.1,label="Comprimento minimo")
+                length_threshold = gr.Slider(0.5,10.,value=3.5,step=0.1,label="Limiar de comprimento")
                 path_precision   = gr.Slider(1,10, value=5,  step=1,   label="Precisao do caminho")
-
             btn = gr.Button("Vetorizar", variant="primary")
-
         with gr.Column():
-            out_svg  = gr.File(label="SVG")
+            out_svg  = gr.File(label="SVG para download")
             out_png  = gr.File(label="PNG vetorizado")
             out_code = gr.Code(label="SVG codigo", language="html", lines=10)
 
     btn.click(
         fn=vectorize_image,
-        inputs=[inp, shadow_strength, n_colors,
-                use_contours, min_area, epsilon,
-                filter_speckle, color_precision,
-                corner_threshold, length_threshold, path_precision],
-        outputs=[out_svg, out_png, out_code]
+        inputs=[inp,shadow_strength,n_colors,
+                filter_speckle,color_precision,
+                corner_threshold,length_threshold,path_precision],
+        outputs=[out_svg,out_png,out_code]
     )
 
 demo.launch()