app.py

import gradio as gr
import vtracer
import os

def convert_svg_to_ai(svg_path):
    """SVG를 AI 파일로 변환"""
    ai_path = svg_path.replace('.svg', '.ai')
    
    # SVG 파일 읽기
    with open(svg_path, 'r') as svg_file:
        svg_content = svg_file.read()
    
    # AI 파일 헤더
    ai_header = """%!PS-Adobe-3.0
%%Creator: Adobe Illustrator(TM) SVG Converter
%%AI8_CreatorVersion: 24.0.0
%%For: SVG to AI Converter
%%Title: Generated AI File
%%CreationDate: %(date)s
%%BoundingBox: 0 0 800 600
%%HiResBoundingBox: 0 0 800 600
%%DocumentData: Clean7Bit
%%LanguageLevel: 2
%%DocumentNeededResources: procset Adobe_packedarray 2.0 0
%%+ procset Adobe_cmykcolor 1.1 0
%%+ procset Adobe_cshow 1.1 0
%%+ procset Adobe_customcolor 1.0 0
%%+ procset Adobe_typography_AI5 1.0 1
%%+ procset Adobe_pattern_AI3 1.0 1
%%+ procset Adobe_Illustrator_AI3 1.0 1
%%EndComments
%%BeginProlog
"""
    
    # AI 파일 생성
    with open(ai_path, 'w', encoding='utf-8') as ai_file:
        # 헤더 작성
        ai_file.write(ai_header)
        
        # SVG 내용을 AI 파일에 포함
        ai_file.write("\n%%BeginDocument\n")
        ai_file.write("/SVGContent\n<<\n")
        ai_file.write("/Type /SVG\n")
        ai_file.write("/Version 1.1\n")
        ai_file.write("/Content [\n")
        ai_file.write(svg_content)
        ai_file.write("\n]\n")
        ai_file.write(">>\ndef\n")
        ai_file.write("\n%%EndDocument\n")
        
        # 파일 종료
        ai_file.write("\n%%Trailer\n%%EOF\n")
    
    return ai_path

def convert_to_vector(
    image, 
    save_svg,
    save_ai,
    colormode="color", 
    hierarchical="stacked", 
    mode="spline", 
    filter_speckle=4, 
    color_precision=6, 
    layer_difference=16, 
    corner_threshold=60, 
    length_threshold=4.0, 
    max_iterations=10, 
    splice_threshold=45, 
    path_precision=3
):
    if not (save_svg or save_ai):
        return None, None, None  # 둘 다 선택되지 않은 경우

    if image is None:
        return None, None, None

    input_path = "temp_input.jpg"
    svg_path = "svg_output.svg"
    outputs = []
    preview = None

    try:
        # 입력 이미지를 임시 파일로 저장
        image.save(input_path)

        # VTracer를 사용하여 이미지를 SVG로 변환
        vtracer.convert_image_to_svg_py(
            input_path,
            svg_path,
            colormode=colormode,
            hierarchical=hierarchical,
            mode=mode,
            filter_speckle=int(filter_speckle),
            color_precision=int(color_precision),
            layer_difference=int(layer_difference),
            corner_threshold=int(corner_threshold),
            length_threshold=float(length_threshold),
            max_iterations=int(max_iterations),
            splice_threshold=int(splice_threshold),
            path_precision=int(path_precision)
        )

        # SVG 파일 처리
        if save_svg:
            with open(svg_path, "r", encoding='utf-8') as f:
                svg_content = f.read()
            preview = gr.HTML(f'<svg viewBox="0 0 {image.width} {image.height}">{svg_content}</svg>')
            outputs.append(svg_path)
        
        # AI 파일 처리
        if save_ai:
            ai_path = convert_svg_to_ai(svg_path)
            outputs.append(ai_path)

        if not save_svg:  # SVG가 선택되지 않았다면 임시 파일 삭제
            os.remove(svg_path)

        return preview, outputs[0] if outputs else None, outputs[1] if len(outputs) > 1 else None

    except Exception as e:
        print(f"Error during conversion: {str(e)}")
        return None, None, None

    finally:
        # 임시 파일 정리
        if os.path.exists(input_path):
            try:
                os.remove(input_path)
            except:
                pass

css = """
#col-container {
    margin: 0 auto;
    max-width: 960px;
}
.generate-btn {
    background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%) !important;
    border: none !important;
    color: white !important;
}
.generate-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
}
"""

# Gradio 인터페이스 정의
with gr.Blocks(css=css) as app:
    with gr.Column(elem_id="col-container"):
        gr.HTML("""
                <div style="text-align: center;">
                    <h2>Image to Vector Converter ⚡</h2>
                    <p>Converts raster images (JPG, PNG, WEBP) to vector graphics (SVG, AI).</p>
                </div>
            """)
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Upload Image")
                with gr.Row():
                    save_svg = gr.Checkbox(value=True, label="Save as SVG")
                    save_ai = gr.Checkbox(value=False, label="Save as AI")
                with gr.Accordion("Advanced Settings", open=False):
                    with gr.Accordion("Clustering", open=False):
                        colormode = gr.Radio([("COLOR","color"),("B/W", "binary")], value="color", label="Color Mode", show_label=False)
                        filter_speckle = gr.Slider(0, 128, value=4, step=1, label="Filter Speckle", info="Cleaner")
                        color_precision = gr.Slider(1, 8, value=6, step=1, label="Color Precision", info="More accurate")
                        layer_difference = gr.Slider(0, 128, value=16, step=1, label="Gradient Step", info="Less layers")
                        hierarchical = gr.Radio([("STACKED","stacked"), ("CUTOUT","cutout")], value="stacked", label="Hierarchical Mode",show_label=False)
                    with gr.Accordion("Curve Fitting", open=False):
                        mode = gr.Radio([("SPLINE","spline"),("POLYGON", "polygon"), ("PIXEL","none")], value="spline", label="Mode", show_label=False)
                        corner_threshold = gr.Slider(0, 180, value=60, step=1, label="Corner Threshold", info="Smoother")
                        length_threshold = gr.Slider(3.5, 10, value=4.0, step=0.1, label="Segment Length", info="More coarse")
                        splice_threshold = gr.Slider(0, 180, value=45, step=1, label="Splice Threshold", info="Less accurate")
                        max_iterations = gr.Slider(1, 20, value=10, step=1, label="Max Iterations", visible=False)
                        path_precision = gr.Slider(1, 10, value=3, step=1, label="Path Precision", visible=False)
                        output_text = gr.Textbox(label="Selected Mode", visible=False)
                with gr.Row():
                    clear_button = gr.Button("Clear")
                    convert_button = gr.Button("✨ Convert", variant='primary', elem_classes=["generate-btn"])

            with gr.Column():
                preview = gr.HTML(label="Preview")
                svg_output = gr.File(label="SVG Output", visible=True)
                ai_output = gr.File(label="AI Output", visible=True)

        examples = [
            ["examples/11.jpg", True, False],
            ["examples/02.jpg", True, False],
            ["examples/03.jpg", True, False],
        ]
        
        gr.Examples(
            examples=examples,
            fn=convert_to_vector,
            inputs=[
                image_input,
                save_svg,
                save_ai
            ],
            outputs=[preview, svg_output, ai_output],
            cache_examples=False,
            run_on_click=True
        )

    # 이벤트 핸들러
    def clear_inputs():
        return [None] * 12

    def update_output_visibility(save_svg, save_ai):
        return gr.update(visible=save_svg), gr.update(visible=save_ai)

    # 체크박스 상태에 따른 출력 파일 컴포넌트 표시/숨김
    save_svg.change(
        update_output_visibility,
        inputs=[save_svg, save_ai],
        outputs=[svg_output, ai_output]
    )
    save_ai.change(
        update_output_visibility,
        inputs=[save_svg, save_ai],
        outputs=[svg_output, ai_output]
    )

    # 변환 버튼 클릭 이벤트
    convert_button.click(
        convert_to_vector,
        inputs=[
            image_input,
            save_svg,
            save_ai,
            colormode,
            hierarchical,
            mode,
            filter_speckle,
            color_precision,
            layer_difference,
            corner_threshold,
            length_threshold,
            max_iterations,
            splice_threshold,
            path_precision
        ],
        outputs=[preview, svg_output, ai_output]
    )

    # Clear 버튼 이벤트
    clear_button.click(
        clear_inputs,
        outputs=[
            image_input,
            colormode,
            hierarchical,
            mode,
            filter_speckle,
            color_precision,
            layer_difference,
            corner_threshold,
            length_threshold,
            max_iterations,
            splice_threshold,
            path_precision
        ]
    )

app.launch(debug=True)