fangxuedeyy/Text_To_Haptics

1import gradio as gr2import matplotlib.pyplot as plt3import os4import json5import random6from huggingface_hub import hf_hub_download, list_repo_files7 8def get_haptic_sample():9    repo_id = "GuiminHu/HapticCap"10    try:11        # 1. 扫描仓库12        print("正在扫描仓库文件夹...")13        all_files = list_repo_files(repo_id, repo_type="dataset")14        15        # 筛选文件夹16        signal_files = [f for f in all_files if f.startswith('haptic_signals/') and f.endswith('.wav')]17        json_files = [f for f in all_files if f.startswith('json/') and f.endswith('.json')]18        19        if not signal_files:20            return "错误：在 haptic_signals/ 文件夹下没找到 .wav 文件", None21 22        # 2. 随机抽取一个信号文件23        test_signal = random.choice(signal_files)24        file_name = os.path.basename(test_signal) # 比如 F100_loop_aug0.wav25        26        # 提取核心 ID (假设核心 ID 是下划线分割的第一部分，如 F100)27        core_id = file_name.split('_')[0] 28        29        # 3. 寻找对应的 JSON 描述30        # 匹配策略：寻找文件名包含核心 ID 的 JSON31        target_json = None32        for jf in json_files:33            if core_id in jf:34                target_json = jf35                break36        37        if not target_json:38            # 如果没找到精准匹配，就随便拿一个 JSON 看看结构，或者报错39            return f"找到了信号 {file_name}，但没找到对应的 JSON。核心ID是: {core_id}", None40 41        # 4. 下载并解析42        sig_path = hf_hub_download(repo_id=repo_id, filename=test_signal, repo_type="dataset")43        json_path = hf_hub_download(repo_id=repo_id, filename=target_json, repo_type="dataset")44        45        with open(json_path, 'r', encoding='utf-8') as f:46            meta = json.load(f)47            # 尝试获取描述字段，HapticCap 可能会把描述放在 'caption' 键里48            caption = meta.get('caption', meta.get('description', 'JSON中未找到描述字段'))49 50        # 5. 绘图 (可视化震动信号)51        import librosa52        signal, sr = librosa.load(sig_path, sr=None)53        54        plt.figure(figsize=(12, 4))55        plt.plot(signal, color='#FF5722', linewidth=0.8)56        plt.title(f"Haptic Waveform: {file_name}")57        plt.xlabel("Time Samples")58        plt.ylabel("Intensity")59        plt.grid(True, linestyle='--', alpha=0.6)60        61        plot_path = "waveform.png"62        plt.savefig(plot_path)63        plt.close()64        65        return f"【文件名】: {file_name}\n【匹配JSON】: {target_json}\n【自然语言描述】: {caption}", plot_path66 67    except Exception as e:68        return f"发生错误: {str(e)}", None69 70# 创建 Gradio 界面71with gr.Blocks(theme=gr.themes.Soft()) as demo:72    gr.Markdown("## 🎧 HapticCap 信号浏览器")73    gr.Markdown("从 75GB 的数据集中随机抽取样本，查看自然语言描述与震动波形的对应关系。")74    75    with gr.Row():76        btn = gr.Button("随机抽取样本", variant="primary")77    78    with gr.Column():79        info_box = gr.Textbox(label="数据详情", lines=5)80        plot_img = gr.Image(label="波形预览")81 82    btn.click(get_haptic_sample, outputs=[info_box, plot_img])83 84demo.launch()