Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import subprocess | |
| import uuid | |
| import sys | |
| import argparse | |
| import shutil | |
| import zipfile | |
| from datetime import timedelta | |
| def timestamp_to_srt(seconds): | |
| td = timedelta(seconds=seconds) | |
| total_seconds = int(td.total_seconds()) | |
| micros = td.microseconds | |
| hours, remainder = divmod(total_seconds, 3600) | |
| minutes, seconds = divmod(remainder, 60) | |
| return f"{hours:02}:{minutes:02}:{seconds:02},{micros//1000:03}" | |
| def json_to_srt(json_data): | |
| """ | |
| Converts internal JSON subtitle format to SRT. | |
| If 'words' key is present, generates word-level timestamps (Karaoke/Editing style). | |
| Otherwise, uses segment-level timestamps. | |
| """ | |
| srt_content = "" | |
| counter = 1 | |
| for block in json_data: | |
| # Check if words detail is available for Word-Level SRT | |
| if isinstance(block, dict) and "words" in block and block["words"]: | |
| for word_obj in block["words"]: | |
| start = word_obj.get('start', 0) | |
| end = word_obj.get('end', 0) | |
| text = word_obj.get('word', "") | |
| srt_content += f"{counter}\n" | |
| srt_content += f"{timestamp_to_srt(start)} --> {timestamp_to_srt(end)}\n" | |
| srt_content += f"{text}\n\n" | |
| counter += 1 | |
| else: | |
| # Fallback to segment level | |
| start = 0 | |
| end = 0 | |
| text = "" | |
| if isinstance(block, dict): | |
| start = block.get('start', 0) | |
| end = block.get('end', 0) | |
| text = block.get('text', "") | |
| elif isinstance(block, (list, tuple)) and len(block) >= 3: | |
| start, end, text = block[0], block[1], block[2] | |
| srt_content += f"{counter}\n" | |
| srt_content += f"{timestamp_to_srt(start)} --> {timestamp_to_srt(end)}\n" | |
| srt_content += f"{text}\n\n" | |
| counter += 1 | |
| return srt_content | |
| # ... (rest of the file until export_pack end) | |
| # 8. ZIP IT | |
| zip_filename = f"{export_name}.zip" | |
| zip_path = os.path.join(project_path, zip_filename) | |
| # Create zip from stage_dir (base_name is without extension) | |
| shutil.make_archive(os.path.join(project_path, export_name), 'zip', stage_dir) | |
| print(f"SUCCESS: Export Pack created at {zip_path}") | |
| # Cleanup | |
| try: | |
| shutil.rmtree(stage_dir) | |
| except: pass | |
| return zip_path | |
| def get_video_dims(vid_path): | |
| """Returns (width, height, duration_frames)""" | |
| try: | |
| cmd_w = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width", "-of", "default=noprint_wrappers=1:nokey=1", vid_path] | |
| width = int(subprocess.check_output(cmd_w).decode().strip()) | |
| cmd_h = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=height", "-of", "default=noprint_wrappers=1:nokey=1", vid_path] | |
| height = int(subprocess.check_output(cmd_h).decode().strip()) | |
| cmd_dur = ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", vid_path] | |
| dur_sec = float(subprocess.check_output(cmd_dur).decode().strip()) | |
| # Assume 30fps for calculation if not probed, but probing is better | |
| cmd_fps = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=r_frame_rate", "-of", "default=noprint_wrappers=1:nokey=1", vid_path] | |
| fps_str = subprocess.check_output(cmd_fps).decode().strip() | |
| num, den = map(int, fps_str.split('/')) | |
| fps = num / den if den > 0 else 30.0 | |
| return width, height, int(dur_sec * fps), fps | |
| except Exception as e: | |
| print(f"Error probing video: {e}") | |
| return 1920, 1080, 300, 30.0 | |
| def render_segmented_overlays(ass_path, segments, video_path, output_dir): | |
| """ | |
| Renders segments using a physical transparent PNG canvas to ensure alpha correctness. | |
| """ | |
| width, height, _, fps = get_video_dims(video_path) | |
| ass_path_sanitized = ass_path.replace("\\", "/").replace(":", "\\:") | |
| # Generate Base Canvas (Robust Way) | |
| canvas_png = os.path.join(output_dir, "base_canvas.png") | |
| # FFmpeg create transparent png | |
| subprocess.run([ | |
| "ffmpeg", "-y", "-f", "lavfi", "-i", f"color=c=black@0.0:s={width}x{height}", | |
| "-frames:v", "1", "-c:v", "png", canvas_png | |
| ], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) | |
| overlay_data = [] | |
| print(f"Rendering {len(segments)} subtitle segments (Mode: Canvas Overlay)...") | |
| for i, seg in enumerate(segments): | |
| start = seg.get('start', 0) | |
| end = seg.get('end', 0) | |
| duration = end - start | |
| if duration <= 0: continue | |
| filename = f"caption_{i}.mov" | |
| out_path = os.path.join(output_dir, filename) | |
| # Input is PNG LOOP (Infinite) -> Trim duration -> Apply ASS -> Encode PNG Codec | |
| cmd = [ | |
| "ffmpeg", "-y", | |
| "-loop", "1", "-i", canvas_png, | |
| "-vf", f"format=rgba,setpts=PTS+{start}/TB,ass='{ass_path_sanitized}',setpts=PTS-{start}/TB,format=rgba", | |
| "-t", str(duration), | |
| "-c:v", "png", | |
| "-pix_fmt", "rgba", | |
| "-an", | |
| out_path | |
| ] | |
| try: | |
| subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) | |
| rel_path = os.path.join("captions", filename).replace("\\", "/") | |
| overlay_data.append({ "path": rel_path, "start": start, "end": end, "index": i }) | |
| print(f" [Seg {i}] Rendered {duration:.2f}s") | |
| except subprocess.CalledProcessError as e: | |
| print(f" [Seg {i}] Failed: {e}") | |
| # Cleanup canvas | |
| if os.path.exists(canvas_png): os.remove(canvas_png) | |
| return overlay_data | |
| def create_premiere_xml(project_name, video_path, overlay_segments, duration_frames, width=1080, height=1920, timebase=30, video_file_id=None, audio_file_id=None, scale_value=100.0, face_data=None, source_width=1920, source_height=1080): | |
| """ | |
| Generates a Premiere Pro XML (xmeml version 4) with segmented cuts and overlays. | |
| overlay_segments: List of dicts [{'path', 'start', 'end', 'index'}] | |
| """ | |
| # Generate unique IDs | |
| def get_uid(): return str(uuid.uuid4())[:12] | |
| if not video_file_id: video_file_id = f"file-video-{get_uid()}" | |
| if not audio_file_id: audio_file_id = f"file-audio-{get_uid()}" | |
| sequence_uuid = str(uuid.uuid4()) | |
| # helper for file blocks | |
| def get_file_block(fid, fpath, is_audio_only=False): | |
| audio_blk = "" if is_audio_only else "<audio><samplecharacteristics><depth>16</depth><samplerate>48000</samplerate></samplecharacteristics><channelcount>2</channelcount></audio>" | |
| width_f = source_width | |
| height_f = source_height | |
| return f"""<file id="{fid}"><name>{os.path.basename(fpath)}</name><pathurl>{fpath}</pathurl><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><duration>{duration_frames}</duration><media><video><samplecharacteristics><width>{width_f}</width><height>{height_f}</height><alpha>straight</alpha></samplecharacteristics></video>{audio_blk}</media></file>""" | |
| # --- VIDEO TRACK SEGMENTATION --- | |
| cuts = [] | |
| frame_face_map = {} | |
| if face_data: | |
| s_factor = scale_value / 100.0 | |
| src_cx = source_width / 2.0 | |
| src_cy = source_height / 2.0 | |
| for entry in face_data: | |
| frame_idx = entry.get('frame') | |
| faces = entry.get('faces', []) | |
| if not faces: continue | |
| best_face = max(faces, key=lambda f: (f[2]-f[0]) * (f[3]-f[1])) | |
| cx = (best_face[0] + best_face[2]) / 2.0 | |
| cy = (best_face[1] + best_face[3]) / 2.0 | |
| off_x = cx - src_cx | |
| off_y = cy - src_cy | |
| tgt_h = - (off_x * s_factor) / width | |
| tgt_v = - (off_y * s_factor) / height | |
| frame_face_map[frame_idx] = (tgt_h, tgt_v) | |
| fps_float = float(timebase) | |
| if overlay_segments: | |
| current_frame = 0 | |
| last_cam_center = (0.0, 0.0) | |
| if 0 in frame_face_map: last_cam_center = frame_face_map[0] | |
| sorted_segs = sorted(overlay_segments, key=lambda x: x['start']) | |
| for seg in sorted_segs: | |
| start_f = int(seg['start'] * fps_float) | |
| end_f = int(seg['end'] * fps_float) | |
| if start_f > current_frame: | |
| cuts.append({ | |
| "start": current_frame, | |
| "end": start_f, | |
| "center": last_cam_center | |
| }) | |
| # Determine Mode Center (Avoid Middle Split) | |
| candidates_h = [] | |
| candidates_v = [] | |
| for f_idx in range(start_f, end_f): | |
| if f_idx in frame_face_map: | |
| pos = frame_face_map[f_idx] | |
| candidates_h.append(round(pos[0], 2)) # Round to cluster | |
| candidates_v.append(round(pos[1], 2)) | |
| if candidates_h: | |
| import statistics | |
| try: | |
| # MODE: Pick the most frequent position | |
| # Multi-mode handling: min(multimode) ensures consistency | |
| best_h = min(statistics.multimode(candidates_h)) | |
| best_v = min(statistics.multimode(candidates_v)) | |
| current_cam_center = (best_h, best_v) | |
| except: | |
| current_cam_center = last_cam_center | |
| else: | |
| current_cam_center = last_cam_center | |
| cuts.append({"start": start_f, "end": end_f, "center": current_cam_center}) | |
| last_cam_center = current_cam_center | |
| current_frame = end_f | |
| if current_frame < duration_frames: | |
| cuts.append({"start": current_frame, "end": duration_frames, "center": last_cam_center}) | |
| else: | |
| cuts.append({"start": 0, "end": duration_frames, "center": (0.0, 0.0)}) | |
| video_track_items = "" | |
| for cut in cuts: | |
| seg_start = cut['start'] | |
| seg_end = cut['end'] | |
| c_h, c_v = cut['center'] | |
| if seg_end - seg_start <= 0: continue | |
| seg_id = f"clipitem-video-{get_uid()}" | |
| basic_motion = f"""<filter><effect><name>Basic Motion</name><effectid>basic</effectid><effectcategory>motion</effectcategory><effecttype>motion</effecttype><mediatype>video</mediatype><parameter authoringApp="PremierePro"><parameterid>scale</parameterid><name>Scale</name><value>{scale_value}</value></parameter><parameter authoringApp="PremierePro"><parameterid>center</parameterid><name>Center</name><value><horiz>{c_h:.5f}</horiz><vert>{c_v:.5f}</vert></value></parameter><parameter authoringApp="PremierePro"><parameterid>centerOffset</parameterid><name>Anchor Point</name><value><horiz>-0.5</horiz><vert>-0.5</vert></value></parameter></effect></filter>""" | |
| video_track_items += f"""<clipitem id="{seg_id}"><name>{os.path.basename(video_path)}</name><duration>{duration_frames}</duration><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><start>{seg_start}</start><end>{seg_end}</end><in>{seg_start}</in><out>{seg_end}</out>{get_file_block(video_file_id, video_path)}{basic_motion}</clipitem>""" | |
| # --- OVERLAY TRACK SEGMENTATION (Opus Style) --- | |
| track_overlay_block = "" | |
| if overlay_segments and len(overlay_segments) > 0: | |
| overlay_clips = "" | |
| fps_float = float(timebase) # Assuming roughly match | |
| for seg in overlay_segments: | |
| start_f = int(seg['start'] * fps_float) | |
| end_f = int(seg['end'] * fps_float) | |
| clip_dur = end_f - start_f | |
| if clip_dur <= 0: continue | |
| ov_name = seg['path'] | |
| ov_fid = f"file-ov-{seg['index']}-{get_uid()}" | |
| ov_cid = f"clip-ov-{seg['index']}-{get_uid()}" | |
| # File block specifically for this WebP | |
| file_blk = f"""<file id="{ov_fid}"><name>{os.path.basename(seg['path'])}</name><pathurl>{seg['path']}</pathurl><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><duration>{clip_dur}</duration><media><video><samplecharacteristics><width>{width}</width><height>{height}</height><alpha>straight</alpha></samplecharacteristics></video></media></file>""" | |
| overlay_clips += f"""<clipitem id="{ov_cid}"><name>{os.path.basename(seg['path'])}</name><duration>{clip_dur}</duration><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><start>{start_f}</start><end>{end_f}</end><in>0</in><out>{clip_dur}</out>{file_blk}<compositemode>normal</compositemode></clipitem>""" | |
| track_overlay_block = f"<track>{overlay_clips}</track>" | |
| else: | |
| track_overlay_block = "<track></track>" | |
| # --- ASSEMBLE --- | |
| timecode_block = f"""<timecode><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><string>00:00:00:00</string><frame>0</frame><displayformat>NDF</displayformat></timecode>""" | |
| audio_blk = f"""<track><clipitem id="{audio_file_id}"><name>{os.path.basename(video_path)}</name><duration>{duration_frames}</duration><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><start>0</start><end>{duration_frames}</end>{get_file_block(video_file_id, video_path)}<sourcetrack><mediatype>audio</mediatype><trackindex>1</trackindex></sourcetrack></clipitem></track>""" | |
| return f"""<?xml version="1.0" encoding="UTF-8"?><xmeml version="4"><sequence id="{sequence_uuid}"><name>{project_name}_CutRef</name><duration>{duration_frames}</duration><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate>{timecode_block}<media><video><format><samplecharacteristics><rate><timebase>{timebase}</timebase><ntsc>FALSE</ntsc></rate><width>{width}</width><height>{height}</height><pixelaspectratio>square</pixelaspectratio></samplecharacteristics></format><track>{video_track_items}</track>{track_overlay_block}</video><audio>{audio_blk}</audio></media></sequence></xmeml>""" | |
| def export_pack(project_path, segment_index, output_format="premiere"): | |
| """ | |
| Generates a ZIP Pack for the segment. | |
| """ | |
| print(f"Starting Export Pack for Project: {os.path.basename(project_path)}, Segment: {segment_index}") | |
| # Paths | |
| proj_name = os.path.basename(project_path) | |
| cut_dir = os.path.join(project_path, "cuts") | |
| # 1. IDENTIFY VIDEO FILE | |
| video_file = None | |
| original_scale_file = None | |
| if os.path.exists(cut_dir): | |
| files = os.listdir(cut_dir) | |
| # Search for {index}_..._original_scale.mp4 or similar | |
| prefix_idx = f"{segment_index:03d}_" | |
| for f in files: | |
| if f.startswith(prefix_idx) and (f.endswith(".mp4") or f.endswith(".mov")): | |
| video_file = os.path.join(cut_dir, f) | |
| break | |
| if not video_file: | |
| print(f"Error: No video file found for segment {segment_index} in {cut_dir}") | |
| return | |
| print(f"Selected Video: {video_file}") | |
| # 2. IDENTIFY SUBTITLE FILES | |
| subs_dir = os.path.join(project_path, "subs_ass") | |
| ass_file = None | |
| if os.path.exists(subs_dir): | |
| sub_files = os.listdir(subs_dir) | |
| prefix_idx = f"{segment_index:03d}_" | |
| # Prioritize Clean Processed > Processed > Any | |
| patterns = [ | |
| (lambda f: f.endswith(".ass") and f.startswith(prefix_idx) and "processed" in f and "original" not in f), | |
| (lambda f: f.endswith(".ass") and f.startswith(prefix_idx) and "processed" in f), | |
| (lambda f: f.endswith(".ass") and f.startswith(prefix_idx)) | |
| ] | |
| for p in patterns: | |
| if ass_file: break | |
| for f in sub_files: | |
| if p(f): | |
| ass_file = os.path.join(subs_dir, f) | |
| break | |
| # JSON in 'subs' usually | |
| subs_json_dir = os.path.join(project_path, "subs") | |
| json_file = None | |
| if os.path.exists(subs_json_dir): | |
| sub_files = os.listdir(subs_json_dir) | |
| prefix_idx = f"{segment_index:03d}_" | |
| # Same pattern priority | |
| json_patterns = [ | |
| (lambda f: f.endswith(".json") and f.startswith(prefix_idx) and "processed" in f), | |
| (lambda f: f.endswith(".json") and f.startswith(prefix_idx)) | |
| ] | |
| for p in json_patterns: | |
| if json_file: break | |
| for f in sub_files: | |
| if p(f): | |
| json_file = os.path.join(subs_json_dir, f) | |
| break | |
| # 2.1 IDENTIFY FACE COORDS | |
| final_dir = os.path.join(project_path, "final") | |
| face_data = None | |
| if os.path.exists(final_dir): | |
| final_files = os.listdir(final_dir) | |
| prefix_idx = f"{segment_index:03d}_" | |
| for f in final_files: | |
| if f.startswith(prefix_idx) and f.endswith("_coords.json"): | |
| try: | |
| with open(os.path.join(final_dir, f), 'r') as fd: | |
| face_data = json.load(fd) | |
| print(f"Found Face Coordinates: {f}") | |
| except Exception as e: | |
| print(f"Face coords load error: {e}") | |
| break | |
| # 3. PREPARE STAGING | |
| export_name = f"export_{proj_name}_seg{segment_index}" | |
| stage_dir = os.path.join(project_path, export_name) | |
| if os.path.exists(stage_dir): | |
| try: | |
| shutil.rmtree(stage_dir) | |
| except Exception: | |
| import random | |
| stage_dir += f"_{random.randint(1000,9999)}" | |
| os.makedirs(stage_dir, exist_ok=True) | |
| # 4. COPY VIDEO | |
| dest_video = os.path.join(stage_dir, "video_cut.mp4") | |
| shutil.copy2(video_file, dest_video) | |
| # 5. RENDER OVERLAYS (SEGMENTED) | |
| overlay_segments = [] | |
| if ass_file and json_file: | |
| try: | |
| with open(json_file, 'r', encoding='utf-8') as f: | |
| jdata = json.load(f) | |
| # Extract segment list | |
| jdata_segs = [] | |
| if isinstance(jdata, dict) and "segments" in jdata: | |
| jdata_segs = jdata["segments"] | |
| elif isinstance(jdata, list): | |
| jdata_segs = jdata | |
| if jdata_segs: | |
| # Create 'captions' subfolder for organization | |
| captions_dir = os.path.join(stage_dir, "captions") | |
| os.makedirs(captions_dir, exist_ok=True) | |
| # Render into subfolder | |
| overlay_segments = render_segmented_overlays(ass_file, jdata_segs, video_file, captions_dir) | |
| except Exception as e: | |
| print(f"Error preparing overlay segments: {e}") | |
| else: | |
| print("Missing ASS or JSON for subtitles. Skipping overlays.") | |
| # 6. GENERATE SRT (Standard) | |
| dest_srt = os.path.join(stage_dir, f"{proj_name}_Seg{segment_index}.srt") | |
| if json_file: | |
| try: | |
| with open(json_file, 'r', encoding='utf-8') as f: | |
| jdata_srt = json.load(f) | |
| if isinstance(jdata_srt, dict) and "segments" in jdata_srt: | |
| jdata_srt = jdata_srt["segments"] | |
| srt_content = json_to_srt(jdata_srt) | |
| with open(dest_srt, 'w', encoding='utf-8') as f: | |
| f.write(srt_content) | |
| except Exception: pass | |
| # 7. GENERATE XML | |
| width_src, height_src, frames, fps = get_video_dims(dest_video) | |
| # Validation for resolution mismatch (same as before) | |
| if face_data: | |
| max_x = 0 | |
| for entry in face_data: | |
| for f in entry.get('faces', []): | |
| if len(f) >= 3 and f[2] > max_x: max_x = f[2] | |
| if max_x > width_src: | |
| print(f"Correction: Detecting 4K source based on face coords ({max_x} > {width_src})") | |
| width_src = 3840 | |
| height_src = 2160 | |
| print(f"Using Source Resolution: {width_src}x{height_src}") | |
| xml_content = create_premiere_xml( | |
| project_name=f"{proj_name}_Seg{segment_index}", | |
| video_path="video_cut.mp4", | |
| overlay_segments=overlay_segments, # LIST PASSED HERE | |
| duration_frames=frames, | |
| width=1080, # Target Vertical Width | |
| height=1920, | |
| timebase=int(float(fps) + 0.5), | |
| source_width=width_src, | |
| source_height=height_src | |
| ) | |
| xml_output = os.path.join(stage_dir, "timeline.xml") | |
| with open(xml_output, "w", encoding="utf-8") as f: | |
| f.write(xml_content) | |
| print("Generated Custom Premiere XML (Opus-Style Segments).") | |
| # 8. ZIP IT | |
| zip_path = f"{stage_dir}.zip" | |
| shutil.make_archive(stage_dir, 'zip', stage_dir) | |
| print(f"SUCCESS: Export Pack created at {zip_path}") | |
| # Cleanup | |
| try: | |
| shutil.rmtree(stage_dir) | |
| except: pass | |
| return zip_path | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--project", required=True) | |
| parser.add_argument("--segment", type=int, required=True) | |
| parser.add_argument("--format", default="premiere") | |
| args = parser.parse_args() | |
| export_pack(args.project, args.segment, args.format) | |