# Attempt to import cv2 and install opencv-python-headless if not available
try:
    import cv2
except ImportError:
    import subprocess
    subprocess.check_call(["pip", "install", "opencv-python-headless"])
    import cv2

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

# Function to convert image to sketch
def convert_to_sketch(img, blur_strength, brightness, contrast):
    # Convert PIL Image to numpy array
    img = np.array(img)
    
    # Convert the image to grayscale
    img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    
    # Invert the grayscale image
    img_inverted = cv2.bitwise_not(img_gray)
    
    # Apply Gaussian blur to the inverted image
    img_blur = cv2.GaussianBlur(img_inverted, (blur_strength, blur_strength), sigmaX=0, sigmaY=0)
    
    # Blend the grayscale and blurred inverted images to create the sketch effect
    img_blend = cv2.divide(img_gray, 255 - img_blur, scale=256)
    
    # Create a white background
    white_background = 255 * np.ones_like(img_blend)
    
    # Add the blended image to the white background
    sketch_with_bg = cv2.addWeighted(img_blend, 1, white_background, 1, 0)
    
    # Adjust brightness and contrast
    sketch_with_bg = adjust_brightness_contrast(sketch_with_bg, brightness, contrast)
    
    # Return the adjusted sketch image
    return sketch_with_bg

# Function to adjust brightness and contrast
def adjust_brightness_contrast(img, brightness, contrast):
    # Apply contrast adjustment
    img = cv2.convertScaleAbs(img, alpha=contrast, beta=brightness)
    return img

# Gradio interface function
def sketch_interface(image, blur_strength, brightness, contrast):
    # Convert the input image to a sketch with adjustments
    sketch = convert_to_sketch(image, blur_strength, brightness, contrast)
    
    # Convert numpy array to PIL Image before returning
    output_image = Image.fromarray(sketch)
    
    # Return the processed sketch image for display
    return output_image

# Create Gradio interface
interface = gr.Interface(
    fn=sketch_interface,
    inputs=[
        gr.Image(type="pil", label="Upload Image", image_mode="RGB", shape=(256, 256)),  # Support for various formats
        gr.Slider(1, 51, step=2, label="Blur Strength", value=21),
        gr.Slider(-100, 100, label="Brightness", value=0),
        gr.Slider(1, 3, step=0.1, label="Contrast", value=1)
    ],
    outputs=[gr.Image(type="pil", label="Sketch Output")],  # Output is now just the image
    title="Cartoon to Sketch Converter",
    description="Upload an image (PNG, JPEG, or others) to convert it into a sketch. Adjust the blur strength, brightness, and contrast for different effects."
)

# Launch the Gradio app
interface.launch()