Dwitee Krishna Panda
adding mindmap
ce47ffe
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr
import os
import torch
import json
import uuid
import langdetect
import moviepy.editor as mp
import yt_dlp
import whisper
from graphviz import Digraph
print("Starting the program...")
# Load BART model for summarization
model_path = "facebook/bart-large-cnn"
print(f"Loading model {model_path}...")
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to("cuda" if torch.cuda.is_available() else "cpu")
model.eval()
print("Model loaded.")
def generate_unique_filename(extension):
return f"{uuid.uuid4()}{extension}"
def cleanup_files(*files):
for file in files:
if file and os.path.exists(file):
os.remove(file)
print(f"Removed file: {file}")
def download_youtube_audio(url):
print(f"Downloading audio from YouTube: {url}")
output_path = generate_unique_filename(".wav")
ydl_opts = {
'format': 'bestaudio/best',
'postprocessors': [{
'key': 'FFmpegExtractAudio',
'preferredcodec': 'wav',
}],
'outtmpl': output_path,
'keepvideo': True,
}
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
ydl.download([url])
return output_path
def transcribe_audio(file_path):
print(f"Transcribing with Whisper: {file_path}")
if file_path.endswith(('.mp4', '.avi', '.mov', '.flv')):
print("Extracting audio from video...")
video = mp.VideoFileClip(file_path)
temp_audio = generate_unique_filename(".wav")
video.audio.write_audiofile(temp_audio)
file_path = temp_audio
model = whisper.load_model("large") # use "base", "medium" if slow
result = model.transcribe(file_path)
print("Transcription done.")
return result["text"]
def generate_summary_stream(transcription):
print("Generating Bullet point summary...")
inputs = tokenizer(transcription, return_tensors="pt", max_length=1024, truncation=True, padding="max_length")
inputs = {k: v.to(model.device) for k, v in inputs.items()}
summary_ids = model.generate(inputs['input_ids'], max_length=300, num_beams=4, early_stopping=True)
raw_summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
bullet_summary = "\n".join(f"• {sentence.strip()}" for sentence in raw_summary.split('.') if sentence.strip())
return bullet_summary
def generate_mindmap_from_summary(summary_text):
dot = Digraph(comment='Mind Map')
dot.node('A', 'Summary')
lines = summary_text.split('\n')
for idx, line in enumerate(lines):
node_id = f'B{idx}'
dot.node(node_id, line.replace("• ", "").strip())
dot.edge('A', node_id)
output_path = generate_unique_filename(".png")
dot.render(output_path, format='png', cleanup=True)
return output_path + ".png"
def process_youtube(url):
if not url:
return "No URL", None
try:
audio_file = download_youtube_audio(url)
transcription = transcribe_audio(audio_file)
return transcription, None
except Exception as e:
return f"Error: {e}", None
finally:
cleanup_files(audio_file)
def process_uploaded_video(video_path):
try:
transcription = transcribe_audio(video_path)
return transcription, None
except Exception as e:
return f"Error: {e}", None
# Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("# 🎥 Video Transcription and Summary")
with gr.Tabs():
with gr.TabItem("📤 Upload Video"):
video_input = gr.Video(label="Upload video")
video_button = gr.Button("🚀 Process Video")
with gr.TabItem("🔗 YouTube Link"):
url_input = gr.Textbox(label="YouTube URL")
url_button = gr.Button("🚀 Process URL")
with gr.Row():
transcription_output = gr.Textbox(label="📝 Transcription", lines=10, show_copy_button=True)
summary_output = gr.Textbox(label="📊 Summary Points", lines=10, show_copy_button=True)
mindmap_output = gr.Image(label="🧠 Mind Map")
summary_button = gr.Button("📝 Generate Summary")
mindmap_button = gr.Button("🧠 Generate Mind Map")
def process_video_and_update(video):
if video is None:
return "No video uploaded", "Please upload a video"
transcription, _ = process_uploaded_video(video)
return transcription or "Transcription error", ""
video_button.click(process_video_and_update, inputs=[video_input], outputs=[transcription_output, summary_output])
url_button.click(process_youtube, inputs=[url_input], outputs=[transcription_output, summary_output])
summary_button.click(generate_summary_stream, inputs=[transcription_output], outputs=[summary_output])
mindmap_button.click(generate_mindmap_from_summary, inputs=[summary_output], outputs=[mindmap_output])
demo.launch(share=True)