import gradio as gr
import numpy as np
import cv2
from PIL import Image
import io
import base64

def convert_to_grayscale(input_image):
    """
    画像を白黒（グレースケール）に変換する関数
    """
    # 入力が PIL Image の場合
    if isinstance(input_image, Image.Image):
        # PIL ImageをOpenCVフォーマットに変換
        img_array = np.array(input_image)
        # RGBからBGRに変換（OpenCVはBGRを使用）
        if len(img_array.shape) == 3:
            img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    else:
        # numpy arrayを直接使用
        img_array = input_image
    
    # グレースケールに変換
    gray_image = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
    
    # グレースケール画像をRGBに戻す（表示用）
    rgb_gray = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2RGB)
    
    return Image.fromarray(rgb_gray)

# REST API用の関数
def api_convert_image(image_data):
    """
    REST API経由で受け取った画像データを処理する関数
    """
    try:
        # 受け取ったデータが既にPIL Imageの場合
        if isinstance(image_data, Image.Image):
            image = image_data
        # base64エンコードされた文字列の場合
        elif isinstance(image_data, str) and image_data.startswith(('data:image', 'data:application')):
            # Base64からデコード
            encoded_data = image_data.split(',')[1]
            decoded_data = base64.b64decode(encoded_data)
            image = Image.open(io.BytesIO(decoded_data))
        # バイナリデータの場合
        elif isinstance(image_data, bytes):
            image = Image.open(io.BytesIO(image_data))
        else:
            return {"error": "不正な画像フォーマットです"}
        
        # グレースケール変換
        gray_image = convert_to_grayscale(image)
        
        # 変換した画像をBase64エンコード
        buffered = io.BytesIO()
        gray_image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        return {
            "status": "success",
            "image": f"data:image/png;base64,{img_str}"
        }
    
    except Exception as e:
        return {"error": str(e)}

# Gradio インターフェース
with gr.Blocks() as demo:
    gr.Markdown("# 画像白黒変換 API")
    gr.Markdown("画像をアップロードすると白黒（グレースケール）に変換します。APIとしても利用可能です。")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="元画像")
        with gr.Column():
            output_image = gr.Image(label="変換後の画像")
    
    convert_btn = gr.Button("変換")
    convert_btn.click(convert_to_grayscale, inputs=input_image, outputs=output_image)
    
    gr.Markdown("## API の使用方法")
    gr.Markdown("""
    このSpaceはAPIとして利用できます。以下の例ではcURLを使用しています：
    
    ```bash
    curl -X POST \
         -F "image=@path/to/your/image.jpg" \
         https://your-username-grayscale-converter.hf.space/api/predict
    ```
    
    または、JSONとして画像をBase64エンコードして送信できます：
    
    ```bash
    curl -X POST \
         -H "Content-Type: application/json" \
         -d '{"image": "data:image/jpeg;base64,/9j/4AAQSkZJRg..."}' \
         https://your-username-grayscale-converter.hf.space/api/predict
    ```
    """)

# APIエンドポイントの定義
demo.queue()
demo.launch()