Create app.py

app.py

@@ -0,0 +1,158 @@
+import gradio as gr
+import cv2
+import whisper
+import spacy
+from PIL import Image
+from diffusers import StableDiffusionPipeline #stable model to be updated if used
+import torch
+import logging
+import os
+import io
+
+# Disable WANDB logging and configure logging level
+logging.disable(logging.WARNING)
+os.environ["WANDB_DISABLED"] = "true"
+
+# Load models
+whisper_model = whisper.load_model("base")
+spacy.prefer_gpu()
+spacy_nlp = spacy.load("en_core_web_sm")
+
+#Initialize the model
+stable_diffusion_pipeline = StableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-2",
+    torch_dtype=torch.float16
+).to("cuda" if torch.cuda.is_available() else "cpu")
+
+
+
+def extract_keyframes(video_path, frame_interval=30, num_frames=5):
+    try:
+        cap = cv2.VideoCapture(video_path)
+        frames = []
+        success, frame = cap.read()
+        count = 0
+        while success and count < num_frames:
+            if count % frame_interval == 0:
+                frames.append(frame)
+            success, frame = cap.read()
+            count += 1
+        cap.release()
+        return frames
+    except Exception as e:
+        logging.error("Error extracting keyframes:", exc_info=e)
+        return None
+def test_extract_keyframes():
+    video_path = "video.mp4"
+    frames = extract_keyframes(video_path)
+
+    assert frames is not None, "Keyframe extraction failed"
+    assert len(frames) > 0, "No keyframes extracted"
+    print("Keyframe extraction test passed")
+
+test_extract_keyframes()
+def transcribe_audio(video_path):
+    try:
+        result = whisper_model.transcribe(video_path)
+        return result['text']
+    except Exception as e:
+        logging.error("Error transcribing audio:", exc_info=e)
+        return None
+def test_transcribe_audio():
+    video_path = "video.mp4"
+    transcription = transcribe_audio(video_path)
+
+    assert transcription is not None, "Transcription failed"
+    assert len(transcription) > 0, "Empty transcription"
+    print("Transcription test passed")
+
+test_transcribe_audio()
+def extract_keywords(text):
+    try:
+        if not text or not text.strip():
+            logging.warning("Empty or whitespace-only text: No keywords extracted")
+            return []
+
+        doc = spacy_nlp(text)
+        keywords = [chunk.text for chunk in doc.noun_chunks]
+
+        if not keywords:
+            logging.warning("No keywords extracted from the text")
+
+        return keywords
+    except Exception as e:
+        logging.error("Error extracting keywords:", exc_info=e)
+        return []
+def test_extract_keywords():
+    text = "This is a test text for keyword extraction."
+    keywords = extract_keywords(text)
+
+    assert keywords is not None, "Keyword extraction failed"
+    assert len(keywords) > 0, "No keywords extracted"
+    print("Keyword extraction test passed")
+
+test_extract_keywords()
+def generate_thumbnails(frames, keywords, num_thumbnails=3):
+    try:
+        thumbnails = []
+        for frame in frames:
+            for _ in range(num_thumbnails):
+                prompt = "A visually striking image of " + ", ".join(keywords)
+                generated_image = stable_diffusion_pipeline(prompt, init_image=frame).images[0]
+                thumbnails.append(generated_image)
+        return thumbnails
+    except Exception as e:
+        logging.exception("Error generating thumbnails:", exc_info=e)
+        return None
+def process_video(video):
+    try:
+        # Determine the video path based on the type of input
+        video_path = video.name if hasattr(video, 'name') else video
+
+        # Extract Keyframes
+        frames = extract_keyframes(video_path)
+        if frames is None:
+            return handle_error("Error extracting keyframes. Please check the video file.")
+
+        # Transcribe Audio
+        transcription = transcribe_audio(video_path)
+        if transcription is None:
+            return handle_error("Error transcribing audio. Please check the audio quality.")
+
+        # Extract Keywords
+        keywords = extract_keywords(transcription)
+        if not keywords:
+            return handle_error("Error extracting keywords. Please check the transcription.")
+
+        # Use the first keyword as title, the full transcription as text, and a generic text placement description
+        title = keywords[0] if keywords else "Thumbnail"
+        text = transcription
+        text_placement = "white letter center at bottom, modern and dynamic"
+
+        # Generate Thumbnails
+        thumbnail_images = generate_thumbnails(frames, keywords)
+        if not thumbnail_images:
+            return handle_error("Error generating thumbnails. Please try again later.")
+
+        return thumbnail_images, "Thumbnails generated successfully."
+    except Exception as e:
+        logging.exception("Unexpected error:", exc_info=e)
+        return handle_error("An unexpected error occurred. Please try again later.")
+def handle_error(error_message):
+    # Return a placeholder image and the error message
+    placeholder = Image.new('RGB', (512, 512), color = (255, 0, 0))  # Placeholder image (red square)
+    return [placeholder], error_message
+# Gradio interface
+interface = gr.Interface(
+    fn=process_video,
+    inputs=gr.Video(label="Upload Video"),
+    outputs=[
+        gr.Gallery(label="Generated Thumbnails"),
+        gr.Textbox(label="Status", lines=2, placeholder="Status message will appear here...")
+    ],
+    title="YouTube Thumbnail Generator",
+    description="Upload a video and generate multiple thumbnails using the video content and transcription.",
+    live=True
+)
+
+interface.launch()