import gradio as gr
import numpy as np
from PIL import Image

def halftone_effect(image, shape="circle", dot_size=10):
    # Convert image to grayscale
    img = Image.fromarray(image).convert('L')
    img_array = np.array(img)
    
    # Get image dimensions
    height, width = img_array.shape
    
    # Create output image
    output = np.ones((height, width, 3), dtype=np.uint8) * 255
    
    # Process each pixel
    for y in range(0, height, dot_size):
        for x in range(0, width, dot_size):
            # Get average brightness in this region
            region = img_array[y:y+dot_size, x:x+dot_size]
            if region.size == 0:
                continue
            brightness = np.mean(region) / 255
            
            # Calculate dot size based on brightness (darker = larger)
            size = int(dot_size * (1 - brightness))
            
            # Draw shape
            if shape == "circle":
                for i in range(max(0, y - size), min(height, y + size)):
                    for j in range(max(0, x - size), min(width, x + size)):
                        if (i - y) ** 2 + (j - x) ** 2 <= size ** 2:
                            output[i, j] = [0, 0, 0]
            
            elif shape == "square":
                for i in range(max(0, y - size), min(height, y + size)):
                    for j in range(max(0, x - size), min(width, x + size)):
                        output[i, j] = [0, 0, 0]
            
            elif shape == "diamond":
                for i in range(max(0, y - size), min(height, y + size)):
                    for j in range(max(0, x - size), min(width, x + size)):
                        if abs(i - y) + abs(j - x) <= size:
                            output[i, j] = [0, 0, 0]
    
    return output

# Create Gradio interface
interface = gr.Interface(
    fn=halftone_effect,
    inputs=[
        gr.Image(label="Input Image"),
        gr.Dropdown(
            choices=["circle", "square", "diamond"],
            value="circle",
            label="Dot Shape"
        ),
        gr.Slider(
            minimum=5,
            maximum=20,
            value=10,
            step=1,
            label="Dot Size"
        )
    ],
    outputs=gr.Image(label="Halftone Result"),
    title="Halftone Effect Generator",
    description="Upload an image and create a halftone effect with different shape options!"
)

# Launch the app
interface.launch()