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

def process_image(input_image):
    # Get image size in KB
    img_size_kb = os.path.getsize(input_image) / 1024
    
    # Create output image (300x300 square)
    size = 300
    if img_size_kb > 500:
        # Green square for large images
        color = [0, 255, 0]  # RGB for green
    else:
        # Black square for small images
        color = [0, 0, 0]  # RGB for black
    
    # Create numpy array with the specified color
    array = np.zeros((size, size, 3), dtype=np.uint8)
    array[:, :] = color
    
    return Image.fromarray(array)

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Image Size Checker")
    gr.Markdown("Upload an image. Output will be green if >500KB, black otherwise.")
    
    with gr.Row():
        input_image = gr.Image(type="filepath", label="Upload Image")
        output_image = gr.Image(label="Result", width=300, height=300)
    
    submit_btn = gr.Button("Check Size")
    submit_btn.click(fn=process_image, inputs=input_image, outputs=output_image)

demo.launch()