Upload New File

src/tape_irregularity_classifier/server.py

+import os
+from concurrent import futures
+import grpc
+from grpc import StatusCode
+from rich.console import Console
+import cv2
+import numpy as np
+
+from tape_irregularity_classifier.lib import (
+    CLASSES,
+    get_irregularity_file,
+    load_model,
+    MODEL_PATH,
+    prepare_images_for_prediction,
+    PROB
+)
+
+from mpai_cae_arp.files import File, FileType
+from mpai_cae_arp.types.irregularity import IrregularityFile, Irregularity, IrregularityType
+from mpai_cae_arp.network import arp_pb2_grpc as arp_pb2_grpc
+from mpai_cae_arp.network.arp_pb2 import (
+    JobRequest,
+    JobResponse,
+    Contact,
+    InfoResponse,
+    License,
+)
+
+PORT = os.getenv("PORT") or '50051'
+info = File('config.yml', FileType.YAML).get_content()
+
+
+def error_response(context, status, message):
+    context.set_code(status)
+    context.set_details(message)
+    return JobResponse(status="error", message=message)
+
+
+class TapeIrregularityClassifierServicer(arp_pb2_grpc.AIMServicer):
+
+    def __init__(self, console: Console):
+        self.console = console
+
+    def getInfo(self, request, context) -> InfoResponse:
+        self.console.log('Received request for AIM info')
+
+        context.set_code(StatusCode.OK)
+        context.set_details('Success')
+
+        return InfoResponse(
+            title=info['title'],
+            description=info['description'],
+            version=info['version'],
+            contact=Contact(
+                name=info['contact']['name'],
+                email=info['contact']['email'],
+            ),
+            license=License(
+                name=info['license_info']['name'],
+                url=info['license_info']['url'],
+            )
+        )
+
+    def work(self, request: JobRequest, context):
+
+        self.console.log('Received request for computation')
+        self.console.log(request)
+        
+        working_dir: str = request.working_dir
+        files_name: str = request.files_name
+
+        temp_dir = os.path.join(working_dir, "temp", files_name)
+        video_irreg_2 = os.path.join(temp_dir, "VideoAnalyser_IrregularityFileOutput2.json")
+        audio_irreg_2 = os.path.join(temp_dir, "AudioAnalyser_IrregularityFileOutput2.json")
+
+        tape_irreg_1 = os.path.join(temp_dir,'TapeIrregularityClassifier_IrregularityFileOutput1.json' )
+        tape_irreg_2 = os.path.join(temp_dir,'TapeIrregularityClassifier_IrregularityFileOutput2.json' )
+
+        audio_irreg_2 = IrregularityFile.from_json(File(audio_irreg_2, FileType.JSON).get_content())
+
+        model = load_model(MODEL_PATH)
+
+        # Open Irregularity File from Video Analyser
+        try:
+            irregularity_file = get_irregularity_file(video_irreg_2)
+        except FileNotFoundError:
+            yield error_response(context, StatusCode.NOT_FOUND, f"{video_irreg_2} not found!")
+        
+        # Get Irregularities array
+        input_irregularities = irregularity_file.irregularities
+        yield JobResponse(
+            status="success",
+            message=f"{len(input_irregularities)} Irregularities listed in {video_irreg_2}")
+
+        # The following array will maintain only the Irregularities that can be processed by the model
+        processed_irregularities: list[Irregularity] = []
+        # Array of images corresponding to the Irregularities
+        images = []
+        yield JobResponse(status="success", message="Loading Irregularity Images...")
+        for irr in input_irregularities:
+            img_path = irr.image_URI
+            if os.path.exists(img_path):
+                # The Irregularity Image was found, then save its data
+                processed_irregularities.append(irr)
+                img = cv2.imread(img_path)
+                # All images must be resized to 224x224 due to network requirements
+                resized = cv2.resize(img, (224, 224), interpolation=cv2.INTER_CUBIC)
+                resized.astype('float32')
+                images.append(resized)
+            else:
+                yield JobResponse(
+                    status="warning",
+                    message=f"Irregularity Image not found! IrregularityID: {irr.irregularity_ID}")
+
+        yield JobResponse(
+            status="success",
+            message=f"{len(processed_irregularities)} Irregularities can be processed by the model")
+
+        images = prepare_images_for_prediction(images)
+        prediction = model.predict(images, batch_size=1)
+        # We are interested in the first array of prediction probabilities
+        probabilities = prediction[0]
+        # Get the position where the highest probability is stored
+        # It corresponds to the class index (0, 1, ...)
+        y = np.argmax(probabilities, axis=1)
+
+        predicted_labels = []
+        # The following array will correspond to the Irregularities that will be listed in the JSON
+        output_irregularities = []
+        for i in range(0, len(processed_irregularities)):
+            # Retrieve current Irregularity probability
+            irr_prob = probabilities[i][y[i]] * 100
+            predicted_labels.append(CLASSES[y[i]] + ' - ' + str(irr_prob) + ' %')
+            irr = processed_irregularities[i]
+            # Consider valid Irregularities only those with probability > prob
+            if probabilities[i][y[i]] * 100 > PROB:
+                irr.irregularity_type = IrregularityType(CLASSES[y[i]])
+                output_irregularities.append(irr)
+            else:
+                yield JobResponse(
+                    status="warning",
+                    message=f"Irregularity number {i+1} with probability {irr_prob:.2f} % discarded")
+                # TODO: Delete corresponding image from temp_path
+                # os.remove(irr['ImageURI'])
+
+        yield JobResponse(
+            status="success",
+            message=f"{len(output_irregularities)} Irregularities with probability > {PROB} %")
+
+        # TODO: IrregularityFileOutput1 shall present only Irregularities relevant to Tape Audio Restoration
+        File(tape_irreg_1, FileType.JSON)\
+            .write_content(
+                IrregularityFile(irregularities=output_irregularities)\
+                    .to_json())
+        yield JobResponse(
+            status="success",
+            message=f"{len(output_irregularities)} Irregularities listed in {tape_irreg_1}")
+
+        File(tape_irreg_2, FileType.JSON)\
+            .write_content(
+                IrregularityFile(irregularities=output_irregularities, offset=audio_irreg_2.offset)\
+                    .to_json())
+        yield JobResponse(
+            status="success",
+            message=f"{len(output_irregularities)} Irregularities listed in {tape_irreg_2}")
+
+
+def serve(console):
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+    arp_pb2_grpc.add_AIMServicer_to_server(TapeIrregularityClassifierServicer(console), server)
+    server.add_insecure_port(f'[::]:{PORT}')
+    server.start()
+    server.wait_for_termination()
+
+
+if __name__ == '__main__':
+    console = Console()
+    console.print(f'Server started at localhost:{PORT} :satellite:')
+    serve(console)