# For video we do the preprocessing and select k sharpest frames. # Afterwards scene is constructed import os # Limit native math-library threads before importing numpy/OpenCV/PyCOLMAP-related code. for _name, _value in { "OPENBLAS_NUM_THREADS": "1", "OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", "NUMEXPR_NUM_THREADS": "1", "VECLIB_MAXIMUM_THREADS": "1", "BLIS_NUM_THREADS": "1", }.items(): os.environ.setdefault(_name, _value) import subprocess import sys from pathlib import Path import cv2 import numpy as np from tqdm import tqdm import time import tempfile from moviepy import VideoFileClip from matplotlib import pyplot as plt from PIL import Image import cv2 from tqdm import tqdm WORKDIR = "../outputs/" def get_rotation_moviepy(video_path): clip = VideoFileClip(video_path) rotation = 0 try: displaymatrix = clip.reader.infos['inputs'][0]['streams'][2]['metadata'].get('displaymatrix', '') if 'rotation of' in displaymatrix: angle = float(displaymatrix.strip().split('rotation of')[-1].split('degrees')[0]) rotation = int(angle) % 360 except Exception as e: print(f"No displaymatrix rotation found: {e}") clip.reader.close() #if clip.audio: # clip.audio.reader.close_proc() return rotation def resize_max_side(frame, max_size): h, w = frame.shape[:2] scale = max_size / max(h, w) if scale < 1: frame = cv2.resize(frame, (int(w * scale), int(h * scale))) return frame def read_video_frames(video_input, k=1, max_size=1024): """ Extracts every k-th frame from a video or list of images, resizes to max size, and returns frames as list. Parameters: video_input (str, file-like, or list): Path to video file, file-like object, or list of image files. k (int): Interval for frame extraction (every k-th frame). max_size (int): Maximum size for width or height after resizing. Returns: frames (list): List of resized frames (numpy arrays). """ # Handle list of image files (not single video in a list) if isinstance(video_input, list): # If it's a single video in a list, treat it as video if len(video_input) == 1 and video_input[0].name.endswith(('.mp4', '.avi', '.mov')): video_input = video_input[0] # unwrap single video file else: # Treat as list of images frames = [] for img_file in video_input: img = Image.open(img_file.name).convert("RGB") img.thumbnail((max_size, max_size)) frames.append(np.array(img)[...,::-1]) return frames # Handle file-like or path if hasattr(video_input, 'name'): video_path = video_input.name elif isinstance(video_input, (str, os.PathLike)): video_path = str(video_input) else: raise ValueError("Unsupported video input type. Must be a filepath, file-like object, or list of images.") cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError(f"Error: Could not open video {video_path}.") total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) frame_count = 0 frames = [] with tqdm(total=total_frames // k, desc="Processing Video", unit="frame") as pbar: while True: ret, frame = cap.read() if not ret: break if frame_count % k == 0: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) h, w = frame.shape[:2] scale = max(h, w) / max_size if scale > 1: frame = cv2.resize(frame, (int(w / scale), int(h / scale))) frames.append(frame[...,[2,1,0]]) pbar.update(1) frame_count += 1 cap.release() return frames def resize_max_side(frame, max_size): """ Resizes the frame so that its largest side equals max_size, maintaining aspect ratio. """ height, width = frame.shape[:2] max_dim = max(height, width) if max_dim <= max_size: return frame # No need to resize scale = max_size / max_dim new_width = int(width * scale) new_height = int(height * scale) resized_frame = cv2.resize(frame, (new_width, new_height), interpolation=cv2.INTER_AREA) return resized_frame def variance_of_laplacian(image): # compute the Laplacian of the image and then return the focus # measure, which is simply the variance of the Laplacian return cv2.Laplacian(image, cv2.CV_64F).var() def process_all_frames(IMG_FOLDER = '/scratch/datasets/hq_data/night2_all_frames', to_visualize=False, save_images=True): dict_scores = {} for idx, img_name in tqdm(enumerate(sorted([x for x in os.listdir(IMG_FOLDER) if '.png' in x]))): img = cv2.imread(os.path.join(IMG_FOLDER, img_name))#[250:, 100:] gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) fm = variance_of_laplacian(gray) + \ variance_of_laplacian(cv2.resize(gray, (0,0), fx=0.75, fy=0.75)) + \ variance_of_laplacian(cv2.resize(gray, (0,0), fx=0.5, fy=0.5)) + \ variance_of_laplacian(cv2.resize(gray, (0,0), fx=0.25, fy=0.25)) if to_visualize: plt.figure() plt.title(f"Laplacian score: {fm:.2f}") plt.imshow(img[..., [2,1,0]]) plt.show() dict_scores[idx] = {"idx" : idx, "img_name" : img_name, "score" : fm} if save_images: dict_scores[idx]["img"] = img return dict_scores def select_optimal_frames(scores, k): """ Selects a minimal subset of frames while ensuring no gaps exceed k. Args: scores (list of float): List of scores where index represents frame number. k (int): Maximum allowed gap between selected frames. Returns: list of int: Indices of selected frames. """ n = len(scores) selected = [0, n-1] i = 0 # Start at the first frame while i < n: # Find the best frame to select within the next k frames best_idx = max(range(i, min(i + k + 1, n)), key=lambda x: scores[x], default=None) if best_idx is None: break # No more frames left selected.append(best_idx) i = best_idx + k + 1 # Move forward, ensuring gaps stay within k return sorted(selected) def variance_of_laplacian(image): """ Compute the variance of Laplacian as a focus measure. """ return cv2.Laplacian(image, cv2.CV_64F).var() def preprocess_frames(frames, verbose=False): """ Compute sharpness scores for a list of frames using multi-scale Laplacian variance. Args: frames (list of np.ndarray): List of frames (BGR images). verbose (bool): If True, print scores. Returns: list of float: Sharpness scores for each frame. """ scores = [] for idx, frame in enumerate(tqdm(frames, desc="Scoring frames")): gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) fm = ( variance_of_laplacian(gray) + variance_of_laplacian(cv2.resize(gray, (0, 0), fx=0.75, fy=0.75)) + variance_of_laplacian(cv2.resize(gray, (0, 0), fx=0.5, fy=0.5)) + variance_of_laplacian(cv2.resize(gray, (0, 0), fx=0.25, fy=0.25)) ) if verbose: print(f"Frame {idx}: Sharpness Score = {fm:.2f}") scores.append(fm) return scores def select_optimal_frames(scores, k): """ Selects k frames by splitting into k segments and picking the sharpest frame from each. Args: scores (list of float): List of sharpness scores. k (int): Number of frames to select. Returns: list of int: Indices of selected frames. """ n = len(scores) selected_indices = [] segment_size = n // k for i in range(k): start = i * segment_size end = (i + 1) * segment_size if i < k - 1 else n # Last chunk may be larger segment_scores = scores[start:end] if len(segment_scores) == 0: continue # Safety check if some segment is empty best_in_segment = start + np.argmax(segment_scores) selected_indices.append(best_in_segment) return sorted(selected_indices) def save_frames_to_scene_dir(frames, scene_dir): """ Saves a list of frames into the target scene directory under 'images/' subfolder. Args: frames (list of np.ndarray): List of frames (BGR images) to save. scene_dir (str): Target path where 'images/' subfolder will be created. """ images_dir = os.path.join(scene_dir, "images") os.makedirs(images_dir, exist_ok=True) for idx, frame in enumerate(frames): filename = os.path.join(images_dir, f"{idx:08d}.png") # 00000000.png, 00000001.png, etc. cv2.imwrite(filename, frame) print(f"Saved {len(frames)} frames to {images_dir}") def run_colmap_on_scene(scene_dir): """Run PyCOLMAP in an isolated subprocess. Native OpenBLAS/FAISS crashes inside pycolmap can otherwise terminate the Gradio backend process. The child process keeps the UI server alive and returns a readable error if COLMAP fails. """ start_time = time.time() print(f"Running isolated COLMAP pipeline on all images inside {scene_dir}") worker_path = Path(__file__).resolve().with_name("colmap_worker.py") env = os.environ.copy() for name, value in { "OPENBLAS_NUM_THREADS": "1", "OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", "NUMEXPR_NUM_THREADS": "1", "VECLIB_MAXIMUM_THREADS": "1", "BLIS_NUM_THREADS": "1", "OPENCV_OPENCL_RUNTIME": "disabled", "OPENCV_OPENCL_DEVICE": "disabled", "EDGS_COLMAP_THREADS": os.getenv("EDGS_COLMAP_THREADS", "2"), "EDGS_COLMAP_MAX_IMAGE_SIZE": os.getenv("EDGS_COLMAP_MAX_IMAGE_SIZE", "1024"), "EDGS_COLMAP_MAX_NUM_FEATURES": os.getenv("EDGS_COLMAP_MAX_NUM_FEATURES", "4096"), "EDGS_COLMAP_MAX_NUM_MATCHES": os.getenv("EDGS_COLMAP_MAX_NUM_MATCHES", "16384"), }.items(): env[name] = value cmd = [sys.executable, str(worker_path), str(scene_dir)] result = subprocess.run( cmd, env=env, cwd=str(Path(__file__).resolve().parents[1]), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, timeout=int(os.getenv("EDGS_COLMAP_TIMEOUT_SECONDS", "900")), ) if result.stdout: print(result.stdout, end="" if result.stdout.endswith("\n") else "\n") if result.returncode != 0: raise RuntimeError( "COLMAP preprocessing failed in the isolated worker. " f"Return code: {result.returncode}. " "Try EDGS_COLMAP_THREADS=1, fewer reference views, or a shorter/higher-overlap clip. " "Worker log tail:\n" + result.stdout[-4000:] ) print(f"Total isolated COLMAP pipeline time: {(time.time() - start_time):.2f}s.")