Jungfraujoch/python/jfjoch_grpc2http.py

#  Copyright (2019-2022) Paul Scherrer Institute
#  SPDX-License-Identifier: GPL-3.0-or-later

from io import BytesIO

import grpc
import numpy
from fastapi import FastAPI, Body, HTTPException, Response
from fastapi.responses import RedirectResponse
from fastapi.staticfiles import StaticFiles
from google.protobuf.json_format import Parse, ParseError, MessageToDict
from tifffile import imwrite

import jfjoch_pb2
import jfjoch_pb2_grpc

static_directory = "/home/leonarski_f/nextgendcu/frontend_ui/build"
grpc_addr = "localhost:5232"

app = FastAPI()

app.mount("/frontend", StaticFiles(directory=static_directory), name="static")

MAX_MESSAGE_LENGTH = 64 * 1024 * 1024
channel = grpc.insecure_channel(
    grpc_addr,
    options=[
        ("grpc.max_send_message_length", MAX_MESSAGE_LENGTH),
        ("grpc.max_receive_message_length", MAX_MESSAGE_LENGTH),
    ],
)


@app.get("/", response_class=RedirectResponse)
async def get_root():
    return "/frontend/index.html"


@app.get("/frontend", response_class=RedirectResponse)
async def get_frontend():
    return "/frontend/index.html"


@app.post("/detector/pedestal")
async def pedestal():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.Pedestal(jfjoch_pb2.Empty())
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())
    return {}


@app.post("/detector/initialize")
async def initialize():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.Initialize(jfjoch_pb2.Empty())
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())
    return {}


@app.post("/detector/deactivate")
async def deactivate():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.Deactivate(jfjoch_pb2.Empty())
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())
    return {}


@app.post("/detector/start")
async def start(data: str = Body(...)):
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.Start(Parse(data, jfjoch_pb2.DatasetSettings()))
        return {}
    except ParseError as e:
        return HTTPException(status_code=400, detail="Parser error: " + str(e))
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.post("/detector/stop")
async def stop():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.Stop(jfjoch_pb2.Empty())
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())
    return {}


@app.post("/detector/trigger")
async def stop():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.Trigger(jfjoch_pb2.Empty())
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())
    return {}


@app.get("/detector/status")
async def get_status():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        return MessageToDict(
            stub.GetStatus(jfjoch_pb2.Empty()), including_default_value_fields=True
        )
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get("/detector/settings")
async def get_settings():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        return MessageToDict(
            stub.GetDetectorSettings(jfjoch_pb2.Empty()),
            including_default_value_fields=True,
        )
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.put("/detector/settings")
async def put_settings(data: str = Body(...)):
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.PutDetectorSettings(Parse(data, jfjoch_pb2.DetectorSettings()))
        return {}
    except ParseError as e:
        return HTTPException(status_code=400, detail="Parser error: " + str(e))
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get("/detector/calibration")
async def get_calibration():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        return MessageToDict(
            stub.GetCalibrationStatistics(jfjoch_pb2.Empty()),
            including_default_value_fields=True,
        )
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get("/data_processing/settings")
async def get_settings():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        return MessageToDict(
            stub.GetDataProcessingSettings(jfjoch_pb2.Empty()),
            including_default_value_fields=True,
        )
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.put("/data_processing/settings")
async def put_data_processing_settings(data: str = Body(...)):
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.PutDataProcessingSettings(Parse(data, jfjoch_pb2.DataProcessingSettings()))
        return {}
    except ParseError as e:
        return HTTPException(status_code=400, detail="Parser error: " + str(e))
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.post("/data_processing/plots")
async def get_settings(data: str = Body(...)):
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        return MessageToDict(
            stub.GetPlots(Parse(data, jfjoch_pb2.PlotRequest())), including_default_value_fields=True
        )
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get("/detector/measurement_statistics")
async def get_meas_stats():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        return MessageToDict(
            stub.GetMeasurementStatistics(jfjoch_pb2.Empty()), including_default_value_fields=True
        )
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get("/detector/list")
async def get_detector_list():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        return MessageToDict(
            stub.GetDetectorList(jfjoch_pb2.Empty()), including_default_value_fields=True
        )
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.put("/detector/select")
async def put_detector_selection(data: str = Body(...)):
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        stub.SelectDetector(Parse(data, jfjoch_pb2.DetectorSelection()))
        return {}
    except ParseError as e:
        return HTTPException(status_code=400, detail="Parser error: " + str(e))
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


# https://stackoverflow.com/questions/55873174/how-do-i-return-an-image-in-fastapi
@app.get(
    "/image/preview.tiff",
    responses={200: {"content": {"image/tiff": {}}}},
    response_class=Response,
)
async def get_preview():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        pbuf = stub.GetPreview(jfjoch_pb2.Empty())
        image_array = numpy.frombuffer(pbuf.data, numpy.int16)
        image_array = numpy.where(image_array < 0, -1, image_array)
        image_array = numpy.reshape(image_array, (pbuf.height, pbuf.width))
        b = BytesIO()
        imwrite(b, image_array)
        return Response(content=b.getvalue(), media_type="image/tiff")
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


def calib_to_tiff(im: jfjoch_pb2.Image) -> bytes:
    if im.pixel_depth == 2:
        image_array = numpy.frombuffer(im.data, numpy.uint16)
    else:
        image_array = numpy.frombuffer(im.data, numpy.uint32)
    image_array = numpy.reshape(image_array, (im.height, im.width))
    b = BytesIO()
    imwrite(b, image_array)
    return bytes(b.getvalue())


@app.get(
    "/image/preview_dioptas.tiff",
    responses={200: {"content": {"image/tiff": {}}}},
    response_class=Response,
)
async def get_preview_dioptas():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        pbuf = stub.GetPreview(jfjoch_pb2.Empty())
        if pbuf.pixel_depth == 2:
            image_array = numpy.frombuffer(pbuf.data, numpy.int16)
        else:
            image_array = numpy.frombuffer(pbuf.data, numpy.int32)
        image_array = numpy.where(image_array < 0, 32767, image_array)
        image_array = numpy.reshape(image_array, (pbuf.height, pbuf.width))
        image_array = numpy.flipud(image_array)
        b = BytesIO()
        imwrite(b, image_array)
        return Response(content=b.getvalue(), media_type="image/tiff")
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get(
    "/image/pedestalG0.tiff",
    responses={200: {"content": {"image/tiff": {}}}},
    response_class=Response,
)
async def get_pedestalg0():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        pbuf = stub.GetPedestalG0(jfjoch_pb2.Empty())
        return Response(content=calib_to_tiff(pbuf), media_type="image/tiff")
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get(
    "/image/pedestalG0.tiff",
    responses={200: {"content": {"image/tiff": {}}}},
    response_class=Response,
)
async def get_pedestalg1():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        pbuf = stub.GetPedestalG1(jfjoch_pb2.Empty())
        return Response(content=calib_to_tiff(pbuf), media_type="image/tiff")
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get(
    "/image/pedestalG2.tiff",
    responses={200: {"content": {"image/tiff": {}}}},
    response_class=Response,
)
async def get_pedestalg2():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        pbuf = stub.GetPedestalG2(jfjoch_pb2.Empty())
        return Response(content=calib_to_tiff(pbuf), media_type="image/tiff")
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())


@app.get(
    "/image/mask.tiff",
    responses={200: {"content": {"image/tiff": {}}}},
    response_class=Response,
)
async def get_mask():
    try:
        stub = jfjoch_pb2_grpc.gRPC_JFJochBrokerStub(channel)
        pbuf = stub.GetMask(jfjoch_pb2.Empty())
        return Response(content=calib_to_tiff(pbuf), media_type="image/tiff")
    except grpc.RpcError as e:
        raise HTTPException(status_code=400, detail=e.details())