替换模型，增加精度

app.py

@@ -8,9 +8,10 @@ import cv2
 import gradio as gr
 import numpy as np
 import torch
+import torch.nn as nn
 import uvicorn
 from fastapi import FastAPI, UploadFile, File, HTTPException
-from torchvision.models.video import r3d_18
+from torchvision.models.video import r3d_18, R3D_18_Weights, r2plus1d_18
 
 from fastapi.staticfiles import StaticFiles
 from fastapi.openapi.docs import (
@@ -24,8 +25,8 @@ static_dir = os.path.dirname(os.path.abspath(__file__))
 app.mount("/static", StaticFiles(directory=f"{static_dir}/static"), name="static")
 
 # 加载预训练视频分类模型（示例使用I3D模型）
-model = r3d_18(pretrained=True)
-model.load_state_dict(torch.load('r3d_18-b3b3357e.pth', map_location='cpu'))
+model = r2plus1d_18(weights="R2Plus1D_18_Weights.KINETICS400_V1")
+model.load_state_dict(torch.load('r2plus1d_18-91a641e6.pth', map_location='cpu'))
 
 # model = torch.hub.load('pytorch/vision', 'r3d_18', pretrained=False)
 # model.load_state_dict(torch.load('r3d_18-b3b3357e.pth'))
@@ -159,6 +160,7 @@ with gr.Blocks(css="""
                    """) as demo:
     gr.Markdown("## 视频内容分类系统")
     with gr.Row():
+        with gr.Column():
         # video_input = gr.UploadButton(
         #     label="上传视频文件",
         #     file_types=["video"],  # 限定视频文件类型[6](@ref)
@@ -166,13 +168,12 @@ with gr.Blocks(css="""
         #     variant="primary"  # 设置按钮样式[2](@ref)
         # )
         # video_input = gr.Video(label="上传视频文件", sources=["upload"])
-        video_input = gr.File(
-            file_type=["mp4", "avi"],
-            label="上传视频文件",
-            multiple=False,
-        )
-        file_name_display = gr.Textbox(label="已上传文件", default="", disabled=True)
-        output_label = gr.Label(label="分类结果TOP5", num_top_classes=5)
+            video_input = gr.File(
+                file_types=[".avi"],
+                label="上传数据集(avi格式)"
+            )
+        with gr.Column():
+            output_label = gr.Label(label="分类结果TOP5", num_top_classes=5)
     submit_btn = gr.Button("开始分析")
     submit_btn.click(fn=video_classifier, inputs=video_input, outputs=output_label)