import base64
import gradio as gr
import requests

def encode_image(image_file):
    with open(image_file.name, "rb") as img_file:
        return base64.b64encode(img_file.read()).decode('utf-8')

def send_to_openai(api_key, image_file):
    base64_image = encode_image(image_file)

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }

    payload = {
        "model": "gpt-4-vision-preview",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Answer in only one of the following options - Leaf , Sheath , Question - You are given a picture of Rice Paddy which part of the Paddy Crop is prominently visible , Leaf should be the whole leaf , and Sheath Can be a little part of the Leaf and Should Show the Stem and Maybe Grains "
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"
                        }
                    }
                ]
            }
        ],
        "max_tokens": 300
    }

    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
    # Extract words from the assistant's response
    assistant_response = response.json()['choices'][0]['message']['content']
    words = assistant_response.split('\n')
    checkresponse_lower = [word.lower() for word in words]
    if "leaf" in checkresponse_lower:
        payload = {
            "model": "gpt-4-vision-preview",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Answer in three words only, does the image uploaded have a healthy rice leaf - Yes or No , does the image uploaded have a rice leaf with Major (not small) circular spots - Yes or No , does the image uploaded have a rice leaf have a major yellowish discoloration in some areas (ignore spots) - Yes or No , DO NOT RESPOND IN MORE THAN THREE WORDS and ANSWER WITH COMMA IN THE MIDDLE OF THE WORDS WITH NO FULLSTOP  "
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ],
            "max_tokens": 300
        }

    elif "sheath" in checkresponse_lower:
        payload = {
            "model": "gpt-4-vision-preview",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "ANSWER IN ONLY ONE WORD , does the sheath part of the paddy in the image have sheath rot ANSWER IN YES OR NO , NO FULLSTOP "
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ],
            "max_tokens": 300
        }
    

    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
    assistant_response = response.json()['choices'][0]['message']['content']
    recognition = assistant_response.split('\n')
    result_string = ' '.join(words + recognition)
    return result_string


iface = gr.Interface(
    fn=send_to_openai,
    inputs=["text", "file"],
    outputs=["text"],
    title="Rice Leaf Disease Detection Demo Using GPT-4V(ision)",
    description="Made By Akash Mondal - https://github.com/akash-mondal                              GPT-4Vision Paper - https://cdn.openai.com/papers/GPTV_System_Card.pdf "
)
iface.launch(debug=True)