#!/usr/bin/env python
# coding: utf-8

# In[4]:


import gradio as gr
from fastai.vision.all import *
import PIL
import torchvision.transforms as T


def get_y(f):
    return int(f.parent.name)-1

# 加载模型
learn = load_learner('pig-model.pkl')

def predict_image(img1, img2):
    
    # 将图像转换为 PIL Image 类型
    img1 = Image.fromarray(img1)  
    img2 = Image.fromarray(img2) 
    

    pred1, _, _ = learn.predict(img1)
    pred2, _, _ = learn.predict(img2)
    
    #print("pred1:",pred1,"pred2:",pred2)

    if pred1 == pred2:
        # 返回预测结果
        return {"是否是同一头猪": "这是同一头猪"}
    else:
        # 返回预测结果
        return {"是否是同一头猪": "这不是同一头猪"}
    


# 创建 Gradio 接口
iface = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.inputs.Image(label="Image 1"),
        gr.inputs.Image(label="Image 2")
    ],
    outputs="text",
    title="猪照片识别模型",
    description="检测两张图片是否是同一只猪。"
)

# 启动 Gradio 接口
iface.launch()