client.py 4.44 KB
Newer Older
Matteo's avatar
update  
Matteo committed
1
2
3
4
5
6
7
8
from rich.console import Console
import os

import grpc
from mpai_cae_arp.network import arp_pb2
from mpai_cae_arp.network import arp_pb2_grpc

channels = {
Matteo's avatar
update    
Matteo committed
9
10
    "AudioAnalyser": grpc.insecure_channel("[::]:50051"),
    "VideoAnalyser": grpc.insecure_channel("[::]:50052"),
Matteo's avatar
Matteo committed
11
    "TapeIrregularityClassifier": grpc.insecure_channel("[::]:50053"),
Matteo's avatar
update    
Matteo committed
12
13
    "TapeAudioRestoration": grpc.insecure_channel("[::]:50054"),
    "Packager": grpc.insecure_channel("[::]:50055"),
Matteo's avatar
update  
Matteo committed
14
15
16
17
}

def run(console: Console):

Matteo's avatar
update    
Matteo committed
18
19
    audio_analyser = arp_pb2_grpc.AIMStub(channels["AudioAnalyser"])
    video_analyser = arp_pb2_grpc.AIMStub(channels["VideoAnalyser"])
Matteo's avatar
Matteo committed
20
    tape_irreg_classifier = arp_pb2_grpc.AIMStub(channels["TapeIrregularityClassifier"])
Matteo's avatar
update    
Matteo committed
21
    tape_audio_restoration = arp_pb2_grpc.AIMStub(channels["TapeAudioRestoration"])
Matteo's avatar
update    
Matteo committed
22
    packager = arp_pb2_grpc.AIMStub(channels["Packager"])
Matteo's avatar
update    
Matteo committed
23
24

    request = arp_pb2.InfoRequest()
Matteo's avatar
update    
Matteo committed
25
26
27
    for aim in [audio_analyser, video_analyser, tape_irreg_classifier, tape_audio_restoration, packager]:
        response = aim.getInfo(request)
        console.print("[bold]{}[/], v{}".format(response.title, response.version))
Matteo's avatar
update  
Matteo committed
28

Matteo's avatar
update    
Matteo committed
29
    request = arp_pb2.JobRequest(
Matteo's avatar
update    
Matteo committed
30
31
        working_dir="/data",
        files_name="BERIO052",
Matteo's avatar
update    
Matteo committed
32
33
34
        index=1,
    )

Matteo's avatar
update    
Matteo committed
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
    with console.status("[bold]Computing AudioAnalyser IrregularityFile 1...", spinner="bouncingBall"):
        for result in audio_analyser.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

    request.files_name = "BERIO052.mov"
    with console.status("[bold]Computing VideoAnalyser IrregularityFiles...", spinner="bouncingBall"):
        for result in video_analyser.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

    request.index = 2
    request.files_name = "BERIO052"
    with console.status("[bold]Computing AudioAnalyser IrregularityFile 2...", spinner="bouncingBall"):
        for result in audio_analyser.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

    with console.status("[bold]Computing TapeIrregularityClassifier...", spinner="bouncingBall"):
        for result in tape_irreg_classifier.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)
Matteo's avatar
update    
Matteo committed
77

Matteo's avatar
update    
Matteo committed
78
79
    with console.status("[bold]Computing TapeAudioRestoration...", spinner="bouncingBall"):
        for result in tape_audio_restoration.work(request):
Matteo's avatar
Matteo committed
80
81
82
83
84
85
86
87
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

Matteo's avatar
update    
Matteo committed
88
89
90
91
92
93
94
95
96
97
    with console.status("[bold]Packaging...", spinner="bouncingBall"):
        for result in packager.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

Matteo's avatar
update    
Matteo committed
98
99
    channels["AudioAnalyser"].close()
    channels["VideoAnalyser"].close()
Matteo's avatar
Matteo committed
100
    channels["TapeIrregularityClassifier"].close()
Matteo's avatar
update    
Matteo committed
101
    channels["TapeAudioRestoration"].close()
Matteo's avatar
update    
Matteo committed
102
    channels["Packager"].close()
Matteo's avatar
update  
Matteo committed
103
104
105
106
107
108

    console.print("[bold green]Success![/] :tada:")


if __name__ == '__main__':
    console = Console()
Matteo's avatar
update    
Matteo committed
109
    run(console)