Merge pull request #1 from FengJin1117/fwh_dev

data_handlers/genre.py

@@ -0,0 +1,50 @@
+from .base import MelodyDatasetHandler
+
+
+class Genre(MelodyDatasetHandler):
+    name = "genre"
+
+    def __init__(self, melody_type, *args, **kwargs):
+        import json
+
+        with open("data/genre/word_data_zh.json", "r", encoding="utf-8") as f:
+            song_db_zh = json.load(f)
+        song_db_zh = {f"zh_{song['id']}": song for song in song_db_zh}  # id as major
+        with open("data/genre/word_data_en.json", "r", encoding="utf-8") as f:
+            song_db_en = json.load(f)
+        song_db_en = {f"en_{song['id']}": song for song in song_db_en} # id as major
+        self.song_db = {**song_db_zh, **song_db_en}
+
+    def get_song_ids(self):
+        return list(self.song_db.keys())
+
+    def get_style_keywords(self, song_id):
+        genre = self.song_db[song_id]["genre"]
+        super_genre = self.song_db[song_id]["super-genre"]
+        gender = self.song_db[song_id]["gender"]
+        return (genre, super_genre, gender)
+
+    def get_phrase_length(self, song_id):
+        # Return the number of lyrics (excluding SP/AP) in each phrase of the song
+        song = self.song_db[song_id]
+        note_lyrics = song.get("note_lyrics", [])
+        
+        phrase_lengths = []
+        for phrase in note_lyrics:
+            count = sum(1 for word in phrase if word not in ("SP", "AP"))
+            phrase_lengths.append(count)
+
+        return phrase_lengths
+
+    def iter_song_phrases(self, song_id):
+        segment_id = 1
+        song = self.song_db[song_id]
+        for phrase_score, phrase_lyrics in zip(song["score"], song["note_lyrics"]):
+            segment = {
+                "note_start_times": [n[0] for n in phrase_score],
+                "note_end_times": [n[1] for n in phrase_score],
+                "note_lyrics": [character for character in phrase_lyrics],
+                "note_midi": [n[2] for n in phrase_score],
+            }
+            yield segment
+            segment_id += 1