(output, file_path,
worker_nodes,
gpu_batchsize,
model_name,
frame_limit=-1,
prefetched_batches=4,
framerate=-1)
| 100 | |
| 101 | |
| 102 | def matte_key(output, file_path, |
| 103 | worker_nodes, |
| 104 | gpu_batchsize, |
| 105 | model_name, |
| 106 | frame_limit=-1, |
| 107 | prefetched_batches=4, |
| 108 | framerate=-1): |
| 109 | manager = multiprocessing.Manager() |
| 110 | |
| 111 | results_dict = manager.dict() |
| 112 | frames_dict = manager.dict() |
| 113 | |
| 114 | |
| 115 | info = ffmpeg.probe(file_path) |
| 116 | cmd = [ |
| 117 | "ffprobe", |
| 118 | "-v", |
| 119 | "error", |
| 120 | "-select_streams", |
| 121 | "v:0", |
| 122 | "-count_packets", |
| 123 | "-show_entries", |
| 124 | "stream=nb_read_packets", |
| 125 | "-of", |
| 126 | "csv=p=0", |
| 127 | file_path |
| 128 | ] |
| 129 | framerate_output = sp.check_output(cmd, universal_newlines=True) |
| 130 | |
| 131 | total_frames = int(framerate_output.split(",")[0]) |
| 132 | if frame_limit != -1: |
| 133 | total_frames = min(frame_limit, total_frames) |
| 134 | |
| 135 | video_stream = next((s for s in info["streams"] if s["codec_type"] == "video"), None) |
| 136 | if not video_stream: |
| 137 | raise Exception("Could not find video stream") |
| 138 | |
| 139 | frame_rate_str = video_stream.get("r_frame_rate", "0/0") |
| 140 | if frame_rate_str == "0/0": |
| 141 | raise Exception("Could not detect framerate of video") |
| 142 | |
| 143 | if framerate == -1: |
| 144 | print(F"FRAME RATE DETECTED: {frame_rate_str} (if this looks wrong, override the frame rate)") |
| 145 | framerate_str = frame_rate_str |
| 146 | framerate_value = _parse_frame_rate(frame_rate_str) |
| 147 | else: |
| 148 | framerate_str = str(framerate) |
| 149 | framerate_value = float(framerate) |
| 150 | |
| 151 | print(F"FRAME RATE: {framerate_value} TOTAL FRAMES: {total_frames}") |
| 152 | |
| 153 | p = multiprocessing.Process(target=capture_frames, |
| 154 | args=(file_path, frames_dict, gpu_batchsize * prefetched_batches, total_frames)) |
| 155 | p.start() |
| 156 | |
| 157 | # note I am deliberately not using pool |
| 158 | # we can't trust it to run all the threads concurrently (or at all) |
| 159 | workers = [multiprocessing.Process(target=worker, |
no test coverage detected