Spaces:
No application file
No application file
File size: 2,390 Bytes
6755a2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import os
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset, BatchSampler, Sampler
from moviepy.editor import VideoFileClip
# TODO: 待后续设计、处理
class OpenCVVideoDataset(Dataset):
def __init__(
self,
video_path,
) -> None:
super().__init__()
self.video_path = video_path
self.cap = cv2.VideoCapture(self.video_path)
def __call__(self, idx) -> np.array:
self.cap.set(2, idx)
frame = self.cap.read()
return frame
def close(self):
self.cap.release()
class MoviepyVideoDataset(Dataset):
def __init__(self, video_path, mode="time") -> None:
self.video_path = video_path
self.videoclip = VideoFileClip(video_path)
self.mode = mode
def __call__(self, t):
if self.mode == "int":
t = t / self.videoclip.fps
frame = self.videoclip.get_frame(t)
return frame
def __len__(self):
n_total = self.videoclip.duration * self.videoclip.fps
return n_total
def generate_videoclip_batchsampler(video_info):
sampler = []
fps = video_info["fps"]
for i, clip in enumerate(video_info["slices"]):
time_start = clip["time_start"]
duration = clip["duration"]
n_start = int(time_start * fps)
n_frame = int(duration * fps)
sampler.append(range(n_start, n_frame, 1))
return sampler
class VideoClipBatchSampler(Sampler):
def __init__(self, sampler) -> None:
self.sampler = sampler
def __iter__(self):
return iter(self.sampler)
def iter_videoclip(model, videoinfo):
pass
if __name__ == "__main__":
import json
PROJECT_DIR = os.path.join(os.path.dirname(__file__), "../..")
DATA_DIR = os.path.join(PROJECT_DIR, "data")
TEST_IMG_PATH = os.path.join(DATA_DIR, "KDA_ALLOUT.jpeg")
TEST_VIDEO_PATH = os.path.join(DATA_DIR, "video.mp4")
TEST_VIDEOMAP_PATH = os.path.join(DATA_DIR, "videomap_大鱼海棠.json")
with open(TEST_VIDEOMAP_PATH, "r") as f:
videoinfo = json.load(f)
videoinfo["fps"] = 30
sampler = generate_videoclip_batchsampler(videoinfo)
videoclip_batchsampler = VideoClipBatchSampler(sampler=sampler)
print("videoclip_batchsampler length", videoclip_batchsampler)
for i, batch in enumerate(videoclip_batchsampler):
print(i, batch)
|