NhanNguyen1309/audio-separation-api

1from __future__ import annotations2 3from app.models import VisualGridCellRecord4 5 6def build_visual_grid_from_beats(7    beats: list[float],8    downbeats: list[float],9    duration_sec: float,10    beats_per_bar: int = 4,11    visual_cells_per_bar: int = 412) -> list[VisualGridCellRecord]:13    if not downbeats:14        return []15 16    cells: list[VisualGridCellRecord] = []17    cell_idx = 118 19    for bar_idx in range(len(downbeats)):20        bar_start = downbeats[bar_idx]21        if bar_start >= duration_sec - 0.02:22            continue23        if bar_idx + 1 < len(downbeats):24            bar_end = downbeats[bar_idx + 1]25        else:26            # estimate from last bar duration27            if bar_idx > 0:28                bar_end = bar_start + (downbeats[bar_idx] - downbeats[bar_idx - 1])29            else:30                # fallback estimate using beats if available31                if len(beats) > 1:32                    bar_end = bar_start + (beats[1] - beats[0]) * beats_per_bar33                else:34                    bar_end = bar_start + 2.0  # fallback 2.0s35 36        # Cap bar_end to duration_sec if needed, but make sure we don't end up with negative/zero duration37        bar_end = min(bar_end, duration_sec)38        if bar_end <= bar_start:39            bar_end = bar_start + 0.00140 41        bar_dur = bar_end - bar_start42        cell_dur = bar_dur / visual_cells_per_bar43 44        for c in range(visual_cells_per_bar):45            c_start = bar_start + c * cell_dur46            c_end = bar_start + (c + 1) * cell_dur47 48            # Clip cell times to duration49            c_start = min(c_start, duration_sec)50            c_end = min(c_end, duration_sec)51 52            # Find nearest beat_index in beats53            nearest_beat_idx = None54            if beats:55                nearest_beat_idx = min(range(len(beats)), key=lambda idx: abs(beats[idx] - c_start))56 57            cells.append(VisualGridCellRecord(58                index=cell_idx,59                visual_bar_index=bar_idx + 1,60                cell_in_bar=c + 1,  # type: ignore[arg-type] # 1 | 2 | 3 | 461                start_sec=round(c_start, 4),62                end_sec=round(c_end, 4),63                is_bar_start=(c == 0),64                beat_index=nearest_beat_idx65            ))66            cell_idx += 167 68    return cells69