VehicleClassification/classify-ai-backend/main.py

import asyncio
import datetime
import os
import websockets
import uuid
import config
from loguru import logger
import json
from algo.car_cf import car_classification

logger.add(config.log_path, rotation="50 MB", encoding='utf-8')

CONNECT_TYPE = 'connect'
ERROR_TYPE = 'error'
CLOSE_TYPE = 'close'
XLS_TYPE = 'xls'
DB_TYPE = 'db'
TXT_TYPE = 'txt'
XLSX_TYPE = 'xlsx'
XML_TYPE = 'xml'


class MessageOperation:

    def __init__(self, uid: str):
        self._uid = uid

    async def handle(self, msg: str) -> str:
        try:
            ask = json.loads(msg)
            self.msg_type = ask['msg_type']
            self.msg_type = ask['msg_type']
            self.record_id = ask['recordId']
            self.msg_content = ask['content']
            if self.msg_type == XLS_TYPE:
                self.response_status, self.response_content = car_classification(
                    self.msg_content)
            elif self.msg_type == DB_TYPE:
                self.response_status, self.response_content = car_classification(
                    self.msg_content)
            elif self.msg_type == TXT_TYPE:
                self.response_status, self.response_content = car_classification(
                    self.msg_content)
            elif self.msg_type == XLSX_TYPE:
                self.response_status, self.response_content = car_classification(
                    self.msg_content)
            elif self.msg_type == XML_TYPE:
                self.response_status, self.response_content = car_classification(
                    self.msg_content)
            elif self.msg_type == CLOSE_TYPE:
                self.response_status = True
                self.response_content = "close success"
            else:
                self.response_status = False
                self.response_content = "无法识别"
                # self.record_id = "null"
        except Exception as e:
            self.msg_type = ERROR_TYPE
            self.response_status = False
            self.response_content = str(e)
            # self.record_id = "null"

    def pack(self) -> str:
        return json.dumps({
            'Sender': self._uid,
            'Recipient': 'server',
            'recordId': self.record_id,
            'status': self.response_status,
            'msg_type': self.msg_type,
            'content': self.response_content,
            'time': str(datetime.datetime.now())
        })


async def main():
    uid = str(uuid.uuid4())
    async with websockets.connect(
            f'{config.WS_URI}?uid={uid}&to_uid=server') as websocket:
        pid = os.getpid()
        logger.info(f'uid: {uid}')
        logger.info(f'pid: {pid}')

        greeting = json.dumps({
            'status': True,
            'msg_type': CONNECT_TYPE,
            'pid': pid,
            'content': '',
            'time': str(datetime.datetime.now())
        })
        logger.debug(f">>> {greeting}")
        await websocket.send(greeting)
        await websocket.recv()
        while True:
            ask = await websocket.recv()
            logger.debug(f"<<< {ask}")
            msg_operation = MessageOperation(uid)
            await msg_operation.handle(ask)
            await websocket.send(msg_operation.pack())
            if msg_operation.msg_type == CLOSE_TYPE:
                break


if __name__ == "__main__":
    logger.info('Starting...')
    asyncio.run(main())
    logger.info('End')