app.py

import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms
import pickle
from resnest.torch import resnest50
from rembg import remove
from PIL import Image
import io
import json
import time
import threading
import concurrent.futures

# 加载类别名称

with open('class_names.pkl', 'rb') as f:
    class_names = pickle.load(f)

# 初始化模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = resnest50(pretrained=False)
model.fc = nn.Sequential(
    nn.Dropout(0.2),
    nn.Linear(model.fc.in_features, len(class_names))
)
model.load_state_dict(torch.load('best_model.pth', map_location=device))
model = model.to(device)
model.eval()

# 预处理流程
preprocess = transforms.Compose([
    transforms.Resize((100, 100)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 创建线程池
executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)


class RealtimeState:
    def __init__(self):
        self.last_result = None
        self.last_update_time = 0
        self.is_processing = False
        self.lock = threading.Lock()


realtime_state = RealtimeState()


def remove_background(img):
    """使用rembg去除背景并添加白色背景"""
    img_byte_arr = io.BytesIO()
    img.save(img_byte_arr, format='PNG')
    img_bytes = img_byte_arr.getvalue()

    removed_bg_bytes = remove(img_bytes)
    removed_bg_img = Image.open(io.BytesIO(removed_bg_bytes)).convert('RGBA')

    white_bg = Image.new('RGBA', removed_bg_img.size, (255, 255, 255, 255))
    combined = Image.alpha_composite(white_bg, removed_bg_img)
    return combined.convert('RGB')


def predict_image(img, remove_bg=False):
    """分类预测主函数"""
    if remove_bg:
        processed_img = remove_background(img)
    else:
        processed_img = img.convert('RGB')

    input_tensor = preprocess(processed_img)
    input_batch = input_tensor.unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(input_batch)

    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    top3_probs, top3_indices = torch.topk(probabilities, 3)

    results = {
        class_names[i]: round(p.item(), 4)
        for p, i in zip(top3_probs, top3_indices)
    }

    best_class = class_names[top3_indices[0]]
    best_conf = top3_probs[0].item() * 100

    with open('prediction_results.txt', 'a') as f:
        f.write(f"Remove BG: {remove_bg}\n")
        f.write(f"Predicted: {best_class} ({best_conf:.2f}%)\n")
        f.write(f"Top 3: {results}\n\n")

    # 添加一个空字符串作为 prediction_id
    prediction_id = ""

    return prediction_id, processed_img, best_class, f"{best_conf:.2f}%", results


def predict_realtime(video_frame, remove_bg):
    """实时预测主函数，结果保留2秒"""
    global realtime_state

    if video_frame is None:
        return None, None, None, None, None

    current_time = time.time()

    # 检查是否有未过期的结果
    with realtime_state.lock:
        if realtime_state.last_result and current_time - realtime_state.last_update_time < 2:
            return realtime_state.last_result

        # 如果正在处理中，返回None
        if realtime_state.is_processing:
            return None, None, None, None, None

        # 标记为正在处理
        realtime_state.is_processing = True

    # 异步处理帧
    def process_frame():
        try:
            result = predict_image(video_frame, remove_bg)
            with realtime_state.lock:
                realtime_state.last_result = result
                realtime_state.last_update_time = time.time()
                realtime_state.is_processing = False
        except Exception as e:
            print(f"处理帧时出错: {e}")
            with realtime_state.lock:
                realtime_state.is_processing = False

    # 提交到线程池处理
    executor.submit(process_frame)

    return None, None, None, None, None


def add_feedback(prediction_id, feedback):
    """模拟将反馈信息保存，实际上不做任何操作"""
    print(f"收到反馈: {feedback} 对预测ID: {prediction_id}")
    return True


def create_interface():
    examples = [
        "r0_0_100.jpg",
        "r0_18_100.jpg",
        "9_100.jpg",
        "5ecc819f1a579f513e0a1500fabb3f0.png",
        "1105.jpg"
    ]

    with gr.Blocks(title="Fruit Classification", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""# 🍎 智能水果识别系统""")

        with gr.Row():
            with gr.Column(scale=3):
                with gr.Group():
                    gr.Markdown("## ⚙️ 处理模式选择")
                    with gr.Row():
                        bg_removal = gr.Checkbox(label="背景去除", value=False, interactive=True)
                with gr.Column():
                    original_image = gr.Image(label="📤 上传图片", type="pil")
                    gr.Examples(examples=examples, inputs=original_image)

                submit_btn = gr.Button("🚀 开始识别", variant="primary")

                gr.Markdown("""## ⚡ 实时识别""")
                camera = gr.Image(label="📷 摄像头捕获", type="pil", streaming=True)

            with gr.Column():
                prediction_id_output = gr.Textbox(label="🔍 预测ID", interactive=False, visible=False)
                processed_image = gr.Image(label="🖼️ 处理后图片", interactive=False)
                best_pred = gr.Textbox(label="🔍 识别结果")
                confidence = gr.Textbox(label="📊 置信度")
                full_results = gr.Label(label="🏆 Top 3 可能结果", num_top_classes=3)

                with gr.Row():
                    feedback_input = gr.Textbox(label="📝 输入反馈信息")
                with gr.Row():
                    feedback_btn = gr.Button("📢 提交反馈", variant="secondary")

        submit_btn.click(
            fn=predict_image,
            inputs=[original_image, bg_removal],
            outputs=[prediction_id_output, processed_image, best_pred, confidence, full_results]
        )

        camera.stream(
            fn=predict_realtime,
            inputs=[camera, bg_removal],
            outputs=[prediction_id_output, processed_image, best_pred, confidence, full_results]
        )

        feedback_btn.click(
            fn=lambda prediction_id, feedback: (
            add_feedback(prediction_id, feedback), "反馈成功！", gr.update(value="")),
            inputs=[prediction_id_output, feedback_input],
            outputs=[prediction_id_output, feedback_input]
        )

    return demo


if __name__ == "__main__":
    interface = create_interface()
    interface.launch(share=True)