Update app.py

app.py

@@ -1,113 +1,70 @@
-# -*- coding: utf-8 -*-
-"""
-鸟类知识科普系统（修正版） by [你的名字]
-ISOM5240 Group Project
-"""
-
-import gradio as gr
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+# 导入必要库
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 from PIL import Image
+import requests
+from io import BytesIO
 import torch
+from transformers import pipeline
+import wikipedia  # 用于获取鸟类百科信息
+from wikipedia.exceptions import DisambiguationError, PageError
 
-# 初始化模型（兼容性优化）
-def init_models():
-    # 鸟类分类模型（保持不变）
-    classifier = pipeline(
-        "image-classification",
-        model="chriamue/bird-species-classifier",
-        device=0 if torch.cuda.is_available() else -1
-    )
+# 1. 鸟类图片识别（使用指定模型）
+def bird_classification(image_url):
+    model_name = "chriamue/bird-species-classifier"
+    feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+    model = AutoModelForImageClassification.from_pretrained(model_name)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
     
-    # 替换为DeepSeek-R1模型（兼容性配置）
-    text_generator = pipeline(
-        "text-generation",
-        model="deepseek-ai/DeepSeek-R1",
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        model_kwargs={
-            "load_in_4bit": True,
-            "trust_remote_code": True  # 必须开启远程代码执行
-        }
-    )
+    # 下载并预处理图片
+    response = requests.get(image_url)
+    img = Image.open(BytesIO(response.content)).convert("RGB")
+    inputs = feature_extractor(img, return_tensors="pt").to(device)
     
-    # 语音合成模型（保持不变）
-    tts = pipeline(
-        "text-to-speech",
-        model="facebook/mms-tts-eng",
-        device=0 if torch.cuda.is_available() else -1
-    )
+    # 模型推理
+    with torch.no_grad():
+        outputs = model(**inputs)
+    probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
     
-    return classifier, text_generator, tts
-
-# 生成儿童友好的鸟类描述（优化Prompt）
-def generate_child_friendly_text(bird_name):
-    PROMPT = f"""以6-12岁儿童能理解的语言介绍{bird_name}：
-1. 用动物拟人化的方式描述特征（例如：穿彩色外套的鸟）
-2. 解释生活习性时结合日常场景（如：像小朋友一样喜欢玩耍）
-3. 包含一个趣味冷知识（例如：飞行距离相当于绕操场XX圈）
-4. 语句长度控制在10-15个英文单词
-5. 使用比喻手法代替专业术语"""
-
-    response = text_generator(
-        PROMPT,
-        max_new_tokens=150,
-        temperature=0.8,
-        top_k=40,
-        do_sample=True
-    )
+    # 获取前1个预测结果
+    predicted_id = torch.argmax(probabilities).item()
+    labels = model.config.id2label
+    bird_species = labels[predicted_id]
+    confidence = round(probabilities[predicted_id].item(), 3)
     
-    # 后处理优化
-    cleaned_text = response[0]['generated_text'].split('\n')[2]
-    return cleaned_text.replace("**", "")  # 去除多余符号
+    return bird_species, confidence
 
-# 主处理流程（增加异常处理）
-def process_image(image):
+# 2. 鸟类信息获取（使用维基百科API）
+def get_bird_info(species_name):
     try:
-        classification = classifier(image)
-        bird_name = classification[0]['label']
-        description = generate_child_friendly_text(bird_name)
-        speech = tts(description, forward_params={"speaker_id": 6})
-        
-        return {
-            "bird_name": bird_name,
-            "description": description,
-            "audio": speech["audio"]
-        }
-    except Exception as e:
-        return f"处理错误: {str(e)}"
+        # 去除可能的多余标签（如模型输出中的括号内容）
+        clean_name = species_name.split("(")[0].strip()
+        # 从维基百科获取摘要（英文转中文）
+        summary = wikipedia.summary(clean_name, sentences=3, auto_suggest=False)
+        return summary
+    except (DisambiguationError, PageError):
+        return "抱歉，未找到该鸟类的详细信息。"
 
-# 初始化模型（增加缓存清理）
-from transformers.utils import cached_file
-cached_file.cache_clear()
-classifier, text_generator, tts = init_models()
+# 3. 文本转语音（使用TTS模型）
+def text_to_speech(text, output_file="bird_info.mp3"):
+    tts = pipeline("text-to-speech", model="tts_models/en_US/tacotron2")
+    speech = tts(text)
+    with open(output_file, "wb") as f:
+        f.write(speech["audio"])
+    return output_file
 
-# 创建Gradio界面（布局优化）
-with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 800px !important}") as demo:
-    gr.Markdown("# 🐦 鸟类知识小课堂（稳定版）")
+# 主函数
+def bird_knowledge_pipeline(image_url):
+    # 1. 鸟类识别
+    species, confidence = bird_classification(image_url)
+    print(f"识别结果：{species}（置信度：{confidence*100:.1f}%）")
     
-    with gr.Row():
-        with gr.Column(scale=2):
-            image_input = gr.Image(type="pil", label="上传鸟类图片", height=300)
-            examples = gr.Examples(
-                examples=["eagle.jpg", "penguin.jpg", "peacock.jpg"],
-                inputs=image_input,
-                label="示例图片"
-            )
-            
-        with gr.Column(scale=3):
-            name_output = gr.Textbox(label="识别到的鸟类", interactive=False)
-            text_output = gr.Textbox(label="趣味知识", lines=4, max_lines=6)
-            audio_output = gr.Audio(label="语音讲解", autoplay=True, visible=True)
+    # 2. 获取详细信息
+    info = get_bird_info(species)
+    print(f"鸟类介绍：\n{info}")
     
-    image_input.change(
-        process_image,
-        inputs=image_input,
-        outputs=[name_output, text_output, audio_output]
-    )
-
-# 部署配置（增加硬件检测）
-if torch.cuda.is_available():
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
-else:
-    print("警告：未检测到GPU，建议在Colab或A10G实例运行")
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    # 3. 生成语音
+    audio_file = text_to_speech(f"这是{species}的介绍：{info}")
+    print(f"语音文件已保存：{audio_file}")
+    
+    return species, info, audio_file