import gradio as gr
import requests
from io import BytesIO

def abnormal_stream(image):
    try:
        byte_io = BytesIO()
        image.save(byte_io, 'png')
        byte_io.seek(0)

        r = requests.post(
            'https://6a051cv20250210-prediction.cognitiveservices.azure.com/customvision/v3.0/Prediction/29f565b7-4710-47a5-8a47-723048ff7ec9/classify/iterations/Iteration2/image',
            headers={
                'Prediction-Key': '8uyKSiqRNbG2JLdMjI8AeOzADtORP3jRh5klqQr0JsJrBBt7x7iPJQQJ99BBACYeBjFXJ3w3AAAIACOGHg4K',
                'Content-Type': 'application/octet-stream',
            },
            data=byte_io,
        )

        if r.status_code != 200:
            return {'확인불가': 1.0, r.status_code: 0.0, r.text: 0.0}

        output_dict = {}

        for item in r.json()['predictions']:
            tag_name = item['tagName']
            probability = item['probability']
            output_dict[tag_name] = probability

        return output_dict

    except Exception as e:
        return {str(e): 1.0}

with gr.Blocks(analytics_enabled=False, title='졸음운전 알리미', head='<link rel="apple-touch-icon" sizes="256x256" href="/pwa_icon" /><meta name="theme-color" content="#0f0f11">') as demo:
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(sources=["webcam"], type="pil")
        with gr.Column():
            output_label = gr.Label()
        dep = input_img.stream(abnormal_stream, [input_img], [output_label])

if __name__ == "__main__":
    demo.launch(favicon_path='./favicon.svg', show_api=False)