from fastapi import WebSocket

class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_json(message)
backend/api/app/connectionManager.py
Racix/Project-AI-Translation
import aiohttp
import os
import pymongo
from app.connectionManager import ConnectionManager
from bson.objectid import ObjectId
from fastapi import FastAPI, HTTPException, UploadFile, status, BackgroundTasks, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import FileResponse
from websockets.exceptions import ConnectionClosedOK

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust this to restrict origins as needed
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

UPLOAD_DIR = os.environ['UPLOAD_DIR']

# Ensure the upload directory exists
os.makedirs(UPLOAD_DIR, exist_ok=True)

# Connect to mongodb
client = pymongo.MongoClient(f"mongodb://{os.environ['MONGO_ADDRESS']}:{os.environ['MONGO_PORT']}/")
db = client["api"]
media_col = db["media"]
analysis_col = db["analysis"]
print(client) #TODO print for debug connection 

# Websocket analysis status manager
analysisManager = ConnectionManager()


@app.get("/media")
async def get_all_media():
    res = list(media_col.find())
    for row in res:
        row['_id'] = str(row['_id'])
    return {"message": res}


@app.post("/media", status_code=status.HTTP_201_CREATED)
async def upload_media(file: UploadFile):

    file_path = os.path.join(UPLOAD_DIR, file.filename)

    # Save the uploaded media locally
    with open(file_path, "wb") as media_file:
        media_file.write(file.file.read())
    
    # Parse data and insert into database
    data = {"name": file.filename, "file_path": file_path}
    res = media_col.insert_one(data)
    
    return {"message": "Media file uploaded successfully", "media_id": str(res.inserted_id)}


@app.get("/media/{media_id}")
async def get_media_by_id(media_id: str):
    media_info = media_col.find_one({"_id": ObjectId(media_id)})
    if media_info is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Media file with id {media_id} was not found")

    media_path = media_info["file_path"]
    if not os.path.exists(media_path):
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"The assisiating media file for id {media_id} was not found")

    return FileResponse(media_info["file_path"], media_type="media/mp4", filename=media_info["name"])


# @app.patch("/media/{media_id}") TODO implement


# @app.delete("/media/{media_id}") TODO implement


@app.post("/media/{media_id}/analysis", status_code=status.HTTP_202_ACCEPTED)
async def start_media_analysis(media_id: str, background_tasks: BackgroundTasks):
    media_info = media_col.find_one({"_id": ObjectId(media_id)})
    if media_info is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Media file not found")

    background_tasks.add_task(analyze, media_id)

    return {"message": "Media file analysis started"}


async def analyze(media_id: str):
    url = f"http://{os.environ['TRANSCRIPTION_ADDRESS']}:{os.environ['API_PORT_GUEST']}/transcribe/{media_id}"
    async with aiohttp.ClientSession() as session:
        res = await session.post(url)
        print("res:", res)


@app.get("/media/{media_id}/analysis")
async def get_media_analysis(media_id: str):
    analysis_info = analysis_col.find_one({"media_id": ObjectId(media_id)}, {"_id": 0, "media_id": 0})
    if analysis_info is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Media file not found")

    return {"message": analysis_info, "media_id": media_id}


# @app.delete("/media/{media_id}/analysis") TODO implement


@app.websocket("/ws/analysis/{media_id}")
async def websocket_endpoint(websocket: WebSocket, media_id: str):
    await analysisManager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_json()
            data["media_id"] = media_id
            await analysisManager.broadcast(data)
    except (WebSocketDisconnect, ConnectionClosedOK):
        pass
    except Exception as e:
        print("Analysis websocket error:", e)
    finally:
        analysisManager.disconnect(websocket)  
        print(f"Client {websocket.client} disconnected")

backend/api/app/main.py
Racix/Project-AI-Translation
import os

def segments_overlap(seg1: list, seg2: list) -> bool:
    """Check if two segments overlap."""
    return seg1["start"] < seg2["start"] + seg2["duration"] and seg1["start"] + seg1["duration"] > seg2["start"]


def align_segments_with_overlap_info(transcript_segments: dict, diarization_segments: dict) -> dict:
    combined_segments = []
    
    for trans_seg in transcript_segments:
        overlapping_dia_segments = [
            dia_seg for dia_seg in diarization_segments 
            if segments_overlap(trans_seg, dia_seg)
        ]
        
        if not overlapping_dia_segments:
            combined_segments.append({
                "text": trans_seg["text"],
                "start": trans_seg["start"],
                "duration": trans_seg["duration"],
                "speaker": "Unknown",
                "overlap_percentage": 0
            })
            continue
        
        if len(overlapping_dia_segments) == 1:
            dia_seg = overlapping_dia_segments[0]
            overlap_start = max(trans_seg["start"], dia_seg["start"])
            overlap_end = min(trans_seg["start"] + trans_seg["duration"], dia_seg["start"] + dia_seg["duration"])
            overlap_duration = overlap_end - overlap_start
            
            overlap_percentage = (overlap_duration / trans_seg["duration"]) * 100
            
            combined_segments.append({
                "text": trans_seg["text"],
                "start": trans_seg["start"],
                "duration": trans_seg["duration"],
                "speaker": dia_seg["speaker"],
                "overlap_percentage": overlap_percentage
            })
        else:
            # If there are multiple overlapping segments, choose the one with the maximum overlap
            overlap_duration = [
                min(seg["start"] + seg["duration"], trans_seg["start"] + trans_seg["duration"]) - 
                max(seg["start"], trans_seg["start"]) 
                for seg in overlapping_dia_segments
            ]
            
            max_index = overlap_duration.index(max(overlap_duration))
            chosen_dia_seg = overlapping_dia_segments[max_index]
            overlap_percentage = (overlap_duration[max_index] / trans_seg["duration"]) * 100
            
            combined_segments.append({
                "text": trans_seg["text"],
                "start": trans_seg["start"],
                "duration": trans_seg["duration"],
                "speaker": chosen_dia_seg["speaker"],
                "overlap_percentage": overlap_percentage
            })

    return combined_segments

# Adjusting the parse_rttm function to read from a file
def parse_rttm_from_file(file_path: str) -> dict:
    file_name, _ = os.path.splitext(os.path.basename(file_path))
    rttm_path = f"/diarization/config/oracle_vad/pred_rttms/{file_name}.rttm" # TODO make this nicer
    segments = []
    with open(rttm_path, 'r', encoding="utf-8") as f:
        for line in f:
            parts = line.strip().split()
            start, duration = float(parts[3]), float(parts[4])
            text = ' '.join(parts[10:]) if len(parts) > 10 else None
            speaker = parts[7] if "speaker" in parts[7] else None
            segments.append({"start": start, "duration": duration, "text": text, "speaker": speaker})
    return segments


backend/diarization/app/combine.py
Racix/Project-AI-Translation
from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels, labels_to_pyannote_object
from nemo.collections.asr.models.msdd_models import NeuralDiarizer
from nemo.collections.asr.models import ClusteringDiarizer
from omegaconf import OmegaConf
from pydub import AudioSegment 
import matplotlib.pyplot as plt
import soundfile as sf
import numpy as np
import subprocess
import librosa
import IPython
import shutil
# import torch
import json
import wget
import sys
import os

CONFIG_DIR = "/diarization/config"
TMP_DIR = "/diarization/tmp"
OUTPUT_DIR = os.path.join(CONFIG_DIR, 'oracle_vad')

os.makedirs(CONFIG_DIR, exist_ok=True)
os.makedirs(TMP_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)


def convert_to_wav(file_path: str, output_path: str):
    """Converts audio file to .wav format."""
    subprocess.call(['ffmpeg', '-hide_banner', '-loglevel warning', '-i', file_path, output_path])


def to_mono(file_path: str):
    """Convert audio file to mono and 16000hz subsample"""
    y, sr = librosa.load(file_path, sr=16000, mono=True)
    sf.write(file_path, y, sr) #TODO maybe not overwrite the original file?


def preprocess(file_path: str) -> str | None:
    file_name = os.path.splitext(os.path.basename(file_path))[0] + ".wav"
    wav_path = os.path.join(TMP_DIR, file_name)
    if file_path.lower().endswith((".mp3", ".mp4")):
        convert_to_wav(file_path, wav_path)
    elif file_path.lower().endswith(".wav"):
        shutil.copyfile(file_path, wav_path)
    else:
        print("wrong file format")
        return None

    to_mono(wav_path)
    return wav_path


def configurations(wav_path: str, domain: str, rttm: str | None, speakers: int) -> OmegaConf:
    # Configuration yaml file
    DOMAIN_TYPE = domain # Can be meeting or telephonic based on domain type of the audio file
    CONFIG_FILE_NAME = f"diar_infer_{DOMAIN_TYPE}.yaml"
    CONFIG_URL = f"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/{CONFIG_FILE_NAME}"
    if not os.path.exists(os.path.join(CONFIG_DIR, CONFIG_FILE_NAME)):
        CONFIG = wget.download(CONFIG_URL, CONFIG_DIR)
    else:
        CONFIG = os.path.join(CONFIG_DIR, CONFIG_FILE_NAME)
    config = OmegaConf.load(CONFIG)
    
    meta = {
        'audio_filepath': wav_path,
        'offset': 0,
        'duration': None,
        'label': 'infer',
        'text': '-',  
        'num_speakers': speakers,
        'rttm_filepath': None, # Sätt till None om vi inte vill ha RTTM!!!!
        'uem_filepath': None
    }

    input_manifest_path = CONFIG_DIR + "/input_manifest.json"
    with open(input_manifest_path,'w') as fp:
        json.dump(meta, fp)
        fp.write('\n')

    #Configure yaml file
    pretrained_speaker_model = "models/titanet-l.nemo"
    pretrained_vad = "models/vad_multilingual_marblenet.nemo"
    pretrained_msdd = "models/diar_msdd_telephonic.nemo"
    config.diarizer.manifest_filepath = input_manifest_path
    # config.device = device
    config.batch_size = 1
    config.diarizer.out_dir = OUTPUT_DIR # Directory to store intermediate files and prediction outputs
    config.diarizer.speaker_embeddings.model_path = pretrained_speaker_model
    config.diarizer.msdd_model.model_path = pretrained_msdd  
    config.diarizer.msdd_model.parameters.sigmoid_threshold = [0.7,1] 
    config.diarizer.msdd_model.parameters.diar_window = 50
    config.diarizer.speaker_embeddings.parameters.window_length_in_sec = [1.5,1.25,1.0,0.75,0.5]
    config.diarizer.speaker_embeddings.parameters.shift_length_in_sec = [0.75,0.625,0.5,0.375,0.1]
    config.diarizer.speaker_embeddings.parameters.multiscale_weights= [1,1,1,1,1]
    config.diarizer.oracle_vad = False # ----> ORACLE VAD Sätt till false om vi inte vill ha RTTM!!!
    config.diarizer.vad.model_path = pretrained_vad
    config.diarizer.clustering.parameters.oracle_num_speakers = False
    return config


def msdd_diarization(config: OmegaConf):
    #Multi-scale model
    oracle_vad_msdd_model = NeuralDiarizer(cfg=config)
    oracle_vad_msdd_model.diarize()


# def cluster_diarization(config: OmegaConf):
#     #Clustering model
#     oracle_vad_clusdiar_model = ClusteringDiarizer(cfg=config)
#     oracle_vad_clusdiar_model.diarize()


def create_diarization(file_path: str, rttm: str | None, speakers: int):
    wav_path = preprocess(file_path)
    config = configurations(wav_path, "telephonic", rttm, speakers)
    #cluster_diarization(config)
    msdd_diarization(config)
backend/diarization/app/diarize.py
Racix/Project-AI-Translation
import app.diarize as di
import app.combine as co
import os
import pymongo
from fastapi import FastAPI, HTTPException, status, BackgroundTasks, Request
from bson.objectid import ObjectId

# Connect to mongodb
client = pymongo.MongoClient(f"mongodb://{os.environ['MONGO_ADDRESS']}:{os.environ['MONGO_PORT']}/")
db = client["api"]
media_col = db["media"]
analysis_col = db["analysis"]

app = FastAPI()


@app.post("/diarize/{media_id}", status_code=status.HTTP_202_ACCEPTED)
async def diarize_media_file(media_id: str, request: Request, background_tasks: BackgroundTasks):
    media_info = media_col.find_one({"_id": ObjectId(media_id)})
    if media_info is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Media file not found")
    
    json = await request.json()
    transcription = json["transcription"]
    if transcription is None:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Transcription data not in request body")

    background_tasks.add_task(diarize, media_id, media_info['file_path'], transcription)

    return {"message": "Diarization of file started"}


async def diarize(media_id: str, file_path: str, transcription: dict):
    analysis_col.delete_one({"media_id": ObjectId(media_id)})
    di.create_diarization(file_path, None, 1) # TODO fix later
    diarization_segments = co.parse_rttm_from_file(file_path)
    transcription['segments'] = co.align_segments_with_overlap_info(transcription['segments'], diarization_segments)
    transcription['media_id'] = ObjectId(media_id)
    analysis_col.insert_one(transcription)
    print("Diarized transcription:", transcription)

    
backend/diarization/app/main.py
Racix/Project-AI-Translation
import aiohttp
import os
import pymongo
import app.transcribe as tr
from bson.objectid import ObjectId
from fastapi import FastAPI, HTTPException, status, BackgroundTasks

# Connect to mongodb
client = pymongo.MongoClient(f"mongodb://{os.environ['MONGO_ADDRESS']}:{os.environ['MONGO_PORT']}/")
db = client["api"]
media_col = db["media"]
print("Mongoclient:", client)

app = FastAPI()


@app.post("/transcribe/{media_id}", status_code=status.HTTP_202_ACCEPTED)
async def transcribe_media_file(media_id: str, background_tasks: BackgroundTasks):
    media_info = media_col.find_one({"_id": ObjectId(media_id)})
    if media_info is None:
        # Notify through websocket transcription failed
        async with aiohttp.ClientSession() as session:
            data = {"status": status.HTTP_404_NOT_FOUND, "message": "Transcription failed. Media file not found."}
            await send_status(data, media_id, session)
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Media file not found")

    background_tasks.add_task(transcribe, media_id, media_info)

    return {"message": "Trancription of file started"}


async def send_status(data: dict, media_id: str, session):
    url = f"ws://{os.environ['API_ADDRESS']}:{os.environ['API_PORT_GUEST']}/ws/analysis/{media_id}"
    async with session.ws_connect(url) as ws:
        await ws.send_json(data)


async def transcribe(media_id: str, media_info):
    async with aiohttp.ClientSession() as session:
        # Notify through websocket transcription started
        data = {"status": status.HTTP_200_OK, "message": "Transcription started..."}
        await send_status(data, media_id, session)

        # Transcribe and send result to diarization model
        trancribe_res = tr.transcribe(media_info['file_path'])
        await send_diarization_req(media_id, trancribe_res)

        # Notify through websocket transcription finished
        data = {"status": status.HTTP_200_OK, "message": "Transcription done."}
        await send_status(data, media_id, session)


async def send_diarization_req(media_id, trancribe_res):
    async with aiohttp.ClientSession() as session:
        url = f"http://{os.environ['DIARIZATION_ADDRESS']}:{os.environ['API_PORT_GUEST']}/diarize/{media_id}"
        res = await session.post(url, json={'transcription': trancribe_res})
        print("res:", res)
backend/transcription/app/main.py
Racix/Project-AI-Translation
from faster_whisper import WhisperModel
import time

model_size = "tiny"
model = WhisperModel(model_size, compute_type="int8")

def transcribe(file_path: str) -> dict:
    start_time = time.time() # TODO only for time print, remove later
    segments, info = model.transcribe(file_path, beam_size=5)
    transcription_segments = []
    for segment in segments:
        transcription_segments.append({
            "text": segment.text,
            "start": segment.start,
            "duration": segment.end - segment.start
        })
    result_dict = {
        "Detected language": info.language,
        "Language probability": info.language_probability * 100,
        "segments": transcription_segments
    }
    # TODO only for time print, remove later
    end_time = time.time() 
    total_time = end_time - start_time
    print("Transcribe time: " + str(total_time) + " seconds")

    return result_dict
backend/transcription/app/transcribe.py
Racix/Project-AI-Translation
from faster_whisper import download_model

download_model("tiny")
backend/transcription/loadModel.py
Racix/Project-AI-Translation