Create main.py

main.py

@@ -0,0 +1,154 @@
+import os
+import time
+import json
+import zipfile
+import shutil
+import cv2
+from pathlib import Path
+from huggingface_hub import HfApi, list_repo_files, hf_hub_download
+
+# ==== CONFIGURATION ====
+HF_TOKEN = os.getenv("HF_TOKEN", "")
+SOURCE_REPO_ID = "factorstudios/movs"
+TARGET_REPO_ID = "factorstudios/movzip"
+
+# Path Configuration
+DOWNLOAD_DIR = "downloads"
+FRAMES_DIR = "frames"
+ZIPS_DIR = "zips"
+
+# Processing Parameters
+TARGET_FPS = 10
+STATE_FILE = "processing_state.json"
+
+# Ensure directories exist
+for d in [DOWNLOAD_DIR, FRAMES_DIR, ZIPS_DIR]:
+    os.makedirs(d, exist_ok=True)
+
+api = HfApi(token=HF_TOKEN)
+
+def log(message):
+    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {message}")
+
+def load_state():
+    if os.path.exists(STATE_FILE):
+        with open(STATE_FILE, 'r') as f:
+            return json.load(f)
+    return {"processed_files": []}
+
+def save_state(state):
+    with open(STATE_FILE, 'w') as f:
+        json.dump(state, f, indent=2)
+
+def extract_frames(video_path, output_dir, fps=TARGET_FPS):
+    """Extract frames from video at specified FPS"""
+    os.makedirs(output_dir, exist_ok=True)
+    cap = cv2.VideoCapture(str(video_path))
+    if not cap.isOpened():
+        log(f"Error: Could not open video {video_path}")
+        return 0
+    
+    video_fps = cap.get(cv2.CAP_PROP_FPS)
+    if not video_fps or video_fps <= 0:
+        video_fps = 30
+    
+    frame_interval = max(1, int(round(video_fps / fps)))
+    
+    frame_idx = 0
+    saved_count = 0
+    
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+            
+        if frame_idx % frame_interval == 0:
+            saved_count += 1
+            frame_name = f"{saved_count:06d}.jpg"
+            cv2.imwrite(os.path.join(output_dir, frame_name), frame, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
+            
+        frame_idx += 1
+        
+    cap.release()
+    return saved_count
+
+def zip_folder(folder_path, zip_path):
+    """Zip all files in a folder"""
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(folder_path):
+            for file in files:
+                zipf.write(os.path.join(root, file), arcname=file)
+
+def process_movies():
+    state = load_state()
+    log(f"Starting movie processing. Source: {SOURCE_REPO_ID}, Target: {TARGET_REPO_ID}")
+    
+    try:
+        # Get list of files in source repo
+        files = list_repo_files(repo_id=SOURCE_REPO_ID, repo_type="dataset")
+        # Filter for video files
+        video_extensions = ('.mp4', '.mkv', '.avi', '.mov', '.webm')
+        videos = [f for f in files if f.lower().endswith(video_extensions)]
+        
+        log(f"Found {len(videos)} videos in source repository.")
+        
+        for video_file in videos:
+            if video_file in state["processed_files"]:
+                log(f"Skipping already processed file: {video_file}")
+                continue
+                
+            log(f"Processing: {video_file}")
+            
+            # 1. Download
+            local_video_path = hf_hub_download(
+                repo_id=SOURCE_REPO_ID,
+                filename=video_file,
+                repo_type="dataset",
+                local_dir=DOWNLOAD_DIR,
+                token=HF_TOKEN
+            )
+            
+            # 2. Extract Frames
+            video_name = Path(video_file).stem
+            video_frames_dir = os.path.join(FRAMES_DIR, video_name)
+            log(f"Extracting frames for {video_name}...")
+            frame_count = extract_frames(local_video_path, video_frames_dir)
+            log(f"Extracted {frame_count} frames.")
+            
+            if frame_count > 0:
+                # 3. Zip
+                zip_filename = f"{video_name}_frames.zip"
+                zip_path = os.path.join(ZIPS_DIR, zip_filename)
+                log(f"Zipping frames to {zip_filename}...")
+                zip_folder(video_frames_dir, zip_path)
+                
+                # 4. Upload
+                log(f"Uploading {zip_filename} to {TARGET_REPO_ID}...")
+                api.upload_file(
+                    path_or_fileobj=zip_path,
+                    path_in_repo=zip_filename,
+                    repo_id=TARGET_REPO_ID,
+                    repo_type="dataset"
+                )
+                log(f"Successfully uploaded {zip_filename}")
+                
+                # Update state
+                state["processed_files"].append(video_file)
+                save_state(state)
+            
+            # 5. Cleanup
+            log(f"Cleaning up local files for {video_name}...")
+            if os.path.exists(video_frames_dir):
+                shutil.rmtree(video_frames_dir)
+            if os.path.exists(local_video_path):
+                os.remove(local_video_path)
+            if 'zip_path' in locals() and os.path.exists(zip_path):
+                os.remove(zip_path)
+                
+        log("All available videos processed.")
+        
+    except Exception as e:
+        log(f"Fatal error during processing: {str(e)}")
+
+if __name__ == "__main__":
+    process_movies()