|
- """
- load an mp4 file
- load music information and calc beat frame index
- detect person and calc ROI region, resize to the same height, eg. 600
-
- save format:
- name
- frame_list []
- poses []
- fps
- width
- height
- """
-
- import librosa
- import cv2
- import time
- import pickle
- import os
-
- class SongLoader():
- def __init__(self, pose_predictor):
- self.song_list = []
- self.save_root = "./songs"
- self.updateSonglist()
- self.pose_predictor = pose_predictor
- self.cur_song = {}
- self.reset_cur()
-
- def reset_cur(self):
- self.play_idx = 0
- self.cur_pose_idx = -1
- self.cur_active = False
-
- def updateSonglist(self):
- self.song_list = []
- files = os.listdir(self.save_root)
- for f in files:
- if f.endswith('pkl'):
- self.song_list.append(f.replace(".pkl", ""))
-
- def save(self, video_path, video_name="dance1"):
- y, sr = librosa.load(video_path)
- _, beats = librosa.beat.beat_track(y=y, sr=sr)
- beat_times = list(librosa.frames_to_time(beats, sr=sr))
- beat_times.append(beat_times[-1] + 1)
- frameIndex = [round(bt * 25) for bt in beat_times]
-
- capture = cv2.VideoCapture(video_path)
- total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
-
- start_ends = []
- last_idx = 0
- for i in range(len(frameIndex) - 1):
- b = frameIndex[i]
- _start = b - 10
-
- if last_idx == 0:
- _start = 0
- elif b - last_idx < 10:
- _start = (b + last_idx) / 2
-
- nb = frameIndex[i + 1]
- _end = b + 10
- if nb - b < 10:
- _end = (b + nb) / 2
-
- if _end > total_frames:
- _end = total_frames
-
- start_ends.append((int(_start), int(_end)))
- last_idx = _end
-
- data = {
- "path": video_path,
- "fps": capture.get(cv2.CAP_PROP_FPS),
- "width": int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)),
- "height": int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)),
- "beats": frameIndex,
- "poses": [],
- "pose_imgs": [],
- "start_ends": start_ends,
- "count": 0
- }
-
- frame_no = 0
- while True:
- ret, frame = capture.read()
- if not ret:
- break
- frame_no += 1
- data["count"] += 1
- if frame_no not in frameIndex:
- continue
-
- frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
- kps, pose_img = self.pose_predictor.predict(frame, True)
- data["poses"].append(kps)
- data["pose_imgs"].append(pose_img)
-
- if len(data["poses"]) < 3:
- return False
-
- save_path = self.save_root + "/{}.pkl".format(video_name)
- with open(save_path, 'wb') as file:
- pickle.dump(data, file)
- self.updateSonglist()
-
- def load(self, song_name):
- path = self.getPklFile(song_name)
- if not os.path.exists(path):
- print("file does not exist: ", path)
- return None
-
- with open(path, "rb") as dbFile:
- self.cur_song = pickle.load(dbFile)
-
- def remove(self, song_name):
- path = self.getPklFile(song_name)
- if os.path.exists(path):
- os.remove(path)
- self.updateSonglist()
- else:
- print('no such file:%s' % song_name)
-
- def getPklFile(self, song_name):
- return f'{self.save_root}/{song_name}.pkl'
-
- def setPlayIdx(self, is_start=False):
- if is_start:
- self.reset_cur()
- else:
- self.play_idx += 1
-
- # leave interval
- if self.cur_active and self.play_idx == self.cur_song["start_ends"][0][1]:
- self.cur_song["start_ends"].pop(0)
- self.cur_active = False
-
- if len(self.cur_song["start_ends"]) == 0:
- return None
-
- # enter interval
- if self.play_idx == self.cur_song["start_ends"][0][0]:
- self.cur_active = True
- self.cur_pose_idx += 1
- print("pose idx: ", self.cur_pose_idx)
- return self.cur_song["pose_imgs"][self.cur_pose_idx : self.cur_pose_idx + 3]
- return None
-
- def getCurVideo(self):
- return self.cur_song["path"]
-
- def getCurPose(self):
- if not self.cur_active:
- return None, -1
-
- if self.cur_pose_idx >= len(self.cur_song["poses"]):
- return None, -1
-
- return self.cur_song["poses"][self.cur_pose_idx], self.cur_pose_idx
-
-
- if __name__ == '__main__':
- song_load = SongLoader()
- print(song_load.song_list)
- song_load.save("/home/hulei/projects/dance/tiger2.mp4", "dance3")
- print(song_load.song_list)
- data = song_load.load("dance1")
- print(data["path"], len(data["pose_imgs"]))
|