VideoAnalysis/main.py

import os
import sys

import cv2 as cv
import numpy as np
import requests
import torch

import imageWorking
import neuralNetwork
import ontologyWorking

url = 'http://kb.athene.tech/api/1.0/ontology/'
img_path = 'data'
img_size = (1280, 720) # Размер изображения для нормализации.


def analyze_file(uid: str, image_path: str) -> None:
    '''
    Анализирует файл и выводит результат в консоль.
    @param uid: УИД онтологии.
    @param url: Базовый URL сервиса.
    '''
    if not ontologyWorking.is_ontology_exists(uid, url):
        raise Exception(f'Онтология с uid {uid} не существует')
    if not os.path.isfile(image_path):
        raise Exception(f'Изображение {image_path} не существует')
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
    model.names = neuralNetwork.rename_entity(model.names)

    # Распознавание изображения.
    results = model(imageWorking.get_image_as_array(image_path))
    
    # Создание аксиом онтологии на основе результатов распознавания.
    object_properties = list()
    data_properties = list()
    for i, res in enumerate(results.pred):
        results_ndarray = np.array(res)
        request = ontologyWorking.get_request_data(model.names, results_ndarray)
        object_properties += request[0]
        data_properties += request[1]

    # Формирование данных для запроса к сервису работы с онтологиями.
    data = {
        'data':
            {
                'objectPropertyAssertions': object_properties,
                'dataPropertyAssertions': data_properties
            }
    }
    result = {
        'QueryGetNotEmpty': '',
        'QueryGetCheck': '',
        'QueryGetEmpty': ''
    }
    params = '&'.join([f'names={query}' for query in result.keys()])
    
    # Выполнение запроса.
    response = requests.post(url + f'{uid}/query/multi?{params}', json=data).json()
    if response['error']:
        raise Exception(response['error'])
    for query in response['response']:
        result[query] = [resultSQWRL['name']['value'] for resultSQWRL in response['response'][query]['rows']]
    print(f'Запрос выполнен')

    # Вывод результата.
    print()
    print('Результат:')
    if result['QueryGetNotEmpty']:
        print('Аудитория занята')
    elif result['QueryGetCheck']:
        print('Аудиторию необходимо проверить')
    elif result['QueryGetEmpty']:
        print('Аудитория пустая')
    else:
        print('Неизвестное состояние')

    # Вывод изображения.
    cv.imshow('result', results.render()[0][:, :, ::-1])
    cv.waitKey(0)
    cv.destroyAllWindows()

# Точка входа в приложение.
if __name__ == '__main__':
    if len(sys.argv) != 3:
        print(f'Запуск: {sys.argv[0]} <Ontology UID> <Image path>')
        exit(1)
    analyze_file(sys.argv[1], sys.argv[2])