Henry Qu commited on
Commit
e749bbe
·
1 Parent(s): 8015d53

new file: main.py

Browse files
Files changed (1) hide show
  1. main.py +63 -0
main.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import os
3
+ from huggingface_hub import hf_hub_download
4
+ from pathlib import Path
5
+ from transformers import GPT2Config, GPT2LMHeadModel, GPT2Tokenizer
6
+
7
+ config_class, model_class, tokenizer_class = GPT2Config, GPT2LMHeadModel, GPT2Tokenizer
8
+ model = model_class.from_pretrained('gpt2')
9
+ tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
10
+
11
+ def search_index(query):
12
+ # 示例返回,实际中应根据查询来搜索索引
13
+ return "example_uuid"
14
+
15
+
16
+ # 下载视频并返回路径的函数
17
+ def download_video(uuid):
18
+ dataset_name = "quchenyuan/360x_dataset"
19
+ dataset_path = "360_dataset/binocular/"
20
+ video_filename = f"{uuid}.mp4"
21
+
22
+ # 确保存储目录存在
23
+ storage_dir = Path("videos")
24
+ storage_dir.mkdir(exist_ok=True)
25
+
26
+ storage_limit = 40*1024 * 1024 * 1024
27
+ current_storage = sum(f.stat().st_size for f in storage_dir.glob('*') if f.is_file())
28
+ if current_storage + os.path.getsize(video_filename) > storage_limit:
29
+ oldest_file = min(storage_dir.glob('*'), key=os.path.getmtime)
30
+ oldest_file.unlink()
31
+
32
+ downloaded_file_path = hf_hub_download(dataset_name, dataset_path + video_filename)
33
+
34
+ return str(storage_dir / video_filename)
35
+
36
+
37
+ # Gradio 接口函数
38
+ def search_and_show_video(query):
39
+ uuid = search_index(query)
40
+ video_path = download_video(uuid)
41
+ return video_path
42
+
43
+
44
+
45
+ if __name__ == "__main__":
46
+ with gr.Blocks() as demo:
47
+ with gr.Column():
48
+ with gr.Row():
49
+ search_input = gr.Textbox(label="输入查询")
50
+ with gr.Row():
51
+ with gr.Column():
52
+ video_output_1 = gr.Video(label="匹配的视频")
53
+ with gr.Column():
54
+ video_output_2 = gr.Video(label="匹配的视频")
55
+ with gr.Column():
56
+ video_output_3 = gr.Video(label="匹配的视频")
57
+ with gr.Row():
58
+ submit_button = gr.Button(label="搜索")
59
+
60
+ submit_button.click(search_and_show_video, search_input, outputs=[video_output_1, video_output_2, video_output_3])
61
+
62
+ # 运行 Gradio 应用
63
+ demo.launch()