import gradio as gr
import torch
from PIL import Image
from ultralytics import YOLO
import matplotlib.pyplot as plt
import io
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
model = YOLO('Fracture_best.pt')

def predict(img, conf, iou):
    results = model.predict(img, conf=conf, iou=iou)
    name = results[0].names
    cls = results[0].boxes.cls
    boneanomaly = 0
    bonelesion = 0
    fracture = 0
    metal = 0
    periostealreaction = 0
    pronatorsign = 0
    softtissue  = 0
    text = 0
    for i in cls:
        if i == 0:
            boneanomaly += 1
        elif i == 1:
            bonelesion += 1
        elif i == 2:
            fracture += 1
        elif i == 3:
            metal += 1
        elif i == 4:
            periostealreaction += 1
        elif i == 5:
            pronatorsign += 1
        elif i==6:
            softtissue += 1
        elif i==8:
            text += 1 
        # 绘制柱状图
    fig, ax = plt.subplots()
    categories = ['BMLY', 'bonelesion', 'fracture', 'metal', 'PN', 'PTS', 'SFTI', 'text']
    counts = [boneanomaly, bonelesion, fracture, metal, periostealreaction, pronatorsign, softtissue, text]
    ax.bar(categories, counts)
    ax.set_title('Category-Count')
    plt.ylim(0,5)
    ax.set_xlabel('Category')
    ax.set_ylabel('Count')
    # 将图表保存为字节流
    buf = io.BytesIO()
    canvas = FigureCanvas(fig)
    canvas.print_png(buf)
    plt.close(fig)  # 关闭图形，释放资源

    # 将字节流转换为PIL Image
    image_png = Image.open(buf)
    # 绘制并返回结果图片和类别计数图表

    for i, r in enumerate(results):
        # Plot results image
        im_bgr = r.plot()  # BGR-order numpy array
        im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image

        # Show results to screen (in supported environments)
    return im_rgb, image_png


base_conf, base_iou = 0.25, 0.45
title = "基于改进YOLOv8算法的手腕骨折辅助诊断系统"
des = "鼠标点击上传图片即可检测缺陷，可通过鼠标调整预测置信度，还可点击网页最下方示例图片进行预测"
interface = gr.Interface(
    inputs=['image', gr.Slider(maximum=1, minimum=0, value=base_conf), gr.Slider(maximum=1, minimum=0, value=base_iou)],
    outputs=["image", 'image'], fn=predict, title=title, description=des,
    examples=[["example1.jpg", base_conf, base_iou],
              ["example2.jpg", base_conf, base_iou],
              ["example3.jpg", base_conf, base_iou]])
interface.launch()