Update app.py

app.py

@@ -1,45 +1,40 @@
-!pip install moviepy
-
-
-
 import streamlit as st
 from huggingface_hub import HfApi
 import os
 import json
 from datetime import datetime
-from moviepy.editor import VideoFileClip
+import cv2
+import random
 from PIL import Image
 import io
-import random
 
 # Initialize the Hugging Face API with the token
 api = HfApi(token=os.getenv("HF_API_TOKEN"))
 
 def extract_thumbnail(video_path, thumbnail_path):
-    try:
-        # Load the video
-        video = VideoFileClip(video_path)
-        
-        # Get a random time point
-        duration = video.duration
-        random_time = random.uniform(0, duration)
-        
-        # Extract the frame
-        frame = video.get_frame(random_time)
-        
-        # Convert the frame to a PIL Image
-        image = Image.fromarray(frame)
-        
-        # Save the thumbnail
-        image.save(thumbnail_path, "JPEG")
-        
-        # Close the video to free up resources
-        video.close()
-        
-        return True
-    except Exception as e:
-        st.error(f"Failed to extract thumbnail: {str(e)}")
-        return False
+    video = cv2.VideoCapture(video_path)
+    
+    # Get total number of frames
+    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    # Choose a random frame
+    random_frame = random.randint(0, total_frames - 1)
+    
+    # Set the frame position
+    video.set(cv2.CAP_PROP_POS_FRAMES, random_frame)
+    
+    # Read the frame
+    success, frame = video.read()
+    if success:
+        cv2.imwrite(thumbnail_path, frame)
+    
+    video.release()
+    return success
+
+def save_custom_thumbnail(thumbnail_file, thumbnail_path):
+    img = Image.open(thumbnail_file)
+    img.save(thumbnail_path)
+    return True
 
 def generate_metadata(video_name, title, description, uploader, file_location, thumbnail_location):
     return {
@@ -54,7 +49,7 @@ def generate_metadata(video_name, title, description, uploader, file_location, t
         "likes": 0
     }
 
-def upload_video_to_hf(video_file, video_name, title, description, uploader):
+def upload_video_to_hf(video_file, video_name, title, description, uploader, custom_thumbnail=None):
     # Create temp paths
     temp_dir = "temp"
     if not os.path.exists(temp_dir):
@@ -70,53 +65,52 @@ def upload_video_to_hf(video_file, video_name, title, description, uploader):
     with open(video_path, "wb") as f:
         f.write(video_file.read())
     
-    # Extract and save thumbnail
-    thumbnail_extracted = extract_thumbnail(video_path, thumbnail_path)
+    # Handle thumbnail
+    if custom_thumbnail:
+        thumbnail_extracted = save_custom_thumbnail(custom_thumbnail, thumbnail_path)
+    else:
+        thumbnail_extracted = extract_thumbnail(video_path, thumbnail_path)
+    
     if not thumbnail_extracted:
+        st.error("Failed to process thumbnail")
         return None
     
-    try:
-        # Upload the video
-        video_location = f"videos/{video_name}"
-        api.upload_file(
-            path_or_fileobj=video_path,
-            path_in_repo=video_location,
-            repo_id="vericudebuget/ok4231",
-            repo_type="space",
-        )
-        
-        # Upload the thumbnail
-        thumbnail_location = f"thumbnails/{thumbnail_name}"
-        api.upload_file(
-            path_or_fileobj=thumbnail_path,
-            path_in_repo=thumbnail_location,
-            repo_id="vericudebuget/ok4231",
-            repo_type="space",
-        )
-        
-        # Generate and upload metadata JSON
-        metadata = generate_metadata(video_name, title, description, uploader, video_location, thumbnail_location)
-        with open(json_path, "w") as f:
-            json.dump(metadata, f, indent=2)
-        
-        api.upload_file(
-            path_or_fileobj=json_path,
-            path_in_repo=f"metadata/{json_name}",
-            repo_id="vericudebuget/ok4231",
-            repo_type="space",
-        )
-        
-        return metadata
+    # Upload the video
+    video_location = f"videos/{video_name}"
+    api.upload_file(
+        path_or_fileobj=video_path,
+        path_in_repo=video_location,
+        repo_id="vericudebuget/ok4231",
+        repo_type="space",
+    )
     
-    except Exception as e:
-        st.error(f"Failed to upload: {str(e)}")
-        return None
+    # Upload the thumbnail
+    thumbnail_location = f"thumbnails/{thumbnail_name}"
+    api.upload_file(
+        path_or_fileobj=thumbnail_path,
+        path_in_repo=thumbnail_location,
+        repo_id="vericudebuget/ok4231",
+        repo_type="space",
+    )
+    
+    # Generate and upload metadata JSON
+    metadata = generate_metadata(video_name, title, description, uploader, video_location, thumbnail_location)
+    with open(json_path, "w") as f:
+        json.dump(metadata, f, indent=2)
     
-    finally:
-        # Cleanup temp files
-        for file_path in [video_path, thumbnail_path, json_path]:
-            if os.path.exists(file_path):
-                os.remove(file_path)
+    api.upload_file(
+        path_or_fileobj=json_path,
+        path_in_repo=f"metadata/{json_name}",
+        repo_id="vericudebuget/ok4231",
+        repo_type="space",
+    )
+    
+    # Cleanup temp files
+    for file_path in [video_path, thumbnail_path, json_path]:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+    
+    return metadata
 
 # Streamlit app interface
 st.title("Upload your video")
@@ -133,6 +127,9 @@ if uploaded_video:
         description = st.text_area("Description", placeholder="Enter video description")
         uploader = st.text_input("Uploader Name", placeholder="Enter your name")
         
+        # Optional custom thumbnail uploader
+        custom_thumbnail = st.file_uploader("Upload custom thumbnail (optional)", type=["jpg", "jpeg", "png"])
+        
         # Upload button within the form
         submit_button = st.form_submit_button("Upload Video")
         
@@ -146,7 +143,8 @@ if uploaded_video:
                         uploaded_video.name, 
                         title, 
                         description, 
-                        uploader
+                        uploader,
+                        custom_thumbnail
                     )
                     if metadata:
                         st.success("Upload completed successfully!")