Update app.py

app.py

@@ -3,11 +3,33 @@ from huggingface_hub import HfApi
 import os
 import json
 from datetime import datetime
+import cv2
+import random
 
 # Initialize the Hugging Face API with the token
 api = HfApi(token=os.getenv("HF_API_TOKEN"))
 
-def generate_metadata(video_name, title, description, uploader, file_location):
+def extract_thumbnail(video_path, thumbnail_path):
+    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 generate_metadata(video_name, title, description, uploader, file_location, thumbnail_location):
     return {
         "fileName": video_name,
         "title": title,
@@ -15,35 +37,53 @@ def generate_metadata(video_name, title, description, uploader, file_location):
         "uploader": uploader,
         "uploadTimestamp": datetime.now().isoformat(),
         "fileLocation": file_location,
+        "thumbnailLocation": thumbnail_location,
         "views": 0,
         "likes": 0
     }
 
 def upload_video_to_hf(video_file, video_name, title, description, uploader):
-    # Save the video temporarily to a file
-    video_path = os.path.join("temp", video_name)
-    json_name = f"{os.path.splitext(video_name)[0]}-index.json"
-    json_path = os.path.join("temp", json_name)
+    # Create temp paths
+    temp_dir = "temp"
+    if not os.path.exists(temp_dir):
+        os.makedirs(temp_dir)
     
-    # Ensure the temp directory exists
-    if not os.path.exists("temp"):
-        os.makedirs("temp")
+    video_path = os.path.join(temp_dir, video_name)
+    thumbnail_name = f"{os.path.splitext(video_name)[0]}_thumb.jpg"
+    thumbnail_path = os.path.join(temp_dir, thumbnail_name)
+    json_name = f"{os.path.splitext(video_name)[0]}-index.json"
+    json_path = os.path.join(temp_dir, json_name)
     
     # Write the video content to a file
     with open(video_path, "wb") as f:
         f.write(video_file.read())
     
-    # Upload the video to Hugging Face
-    file_location = f"videos/{video_name}"
+    # Extract and save thumbnail
+    thumbnail_extracted = extract_thumbnail(video_path, thumbnail_path)
+    if not thumbnail_extracted:
+        st.error("Failed to extract thumbnail from video")
+        return None
+    
+    # Upload the video
+    video_location = f"videos/{video_name}"
     api.upload_file(
         path_or_fileobj=video_path,
-        path_in_repo=file_location,
+        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, file_location)
+    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)
     
@@ -54,6 +94,11 @@ def upload_video_to_hf(video_file, video_name, title, description, uploader):
         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
@@ -78,7 +123,7 @@ if uploaded_video:
             if not title or not uploader:
                 st.error("Please fill in the title and uploader name.")
             else:
-                with st.spinner("Uploading video and generating metadata..."):
+                with st.spinner("Uploading video, generating thumbnail and metadata..."):
                     metadata = upload_video_to_hf(
                         uploaded_video, 
                         uploaded_video.name, 
@@ -86,8 +131,9 @@ if uploaded_video:
                         description, 
                         uploader
                     )
-                    st.success("Upload completed successfully!")
-                    st.json(metadata)
+                    if metadata:
+                        st.success("Upload completed successfully!")
+                        st.json(metadata)
 
 else:
     st.info("Please upload a video file to begin.")