DroneDetector/NN_server/Models/ensemble_915_v44.py

from torchvision import models
import torch.nn as nn
import matplotlib
import numpy as np
import torch
import cv2
import gc
import io


def _render_plot(values, figsize=(16, 16), dpi=16):
    import matplotlib.pyplot as plt

    fig = plt.figure(figsize=figsize)
    plt.axes(ylim=(-1, 1))
    plt.plot(values, color="black")
    plt.gca().set_axis_off()
    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
    plt.margins(0, 0)

    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=dpi)
    buf.seek(0)
    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
    buf.close()

    img = cv2.imdecode(img_arr, 1)
    if img is None:
        raise RuntimeError("failed to decode plot image")

    plt.clf()
    plt.cla()
    plt.close()
    plt.close(fig)

    return np.asarray(cv2.split(img), dtype=np.float32)


def pre_func_ensemble(data=None, src="", ind_inference=0):
    try:
        import matplotlib.pyplot as plt

        matplotlib.use("Agg")
        plt.ioff()

        real = np.asarray(data[0], dtype=np.float32)
        imag = np.asarray(data[1], dtype=np.float32)
        signal = real + 1j * imag

        img_real = _render_plot(signal.real)
        img_mag = _render_plot(np.abs(signal))

        cv2.destroyAllWindows()
        gc.collect()

        print("Подготовка данных завершена")
        print()
        return [img_real, img_mag]

    except Exception as exc:
        print(str(exc))
        return None


def build_func_ensemble(file_model="", file_config="", num_classes=None):
    try:
        import matplotlib.pyplot as plt

        matplotlib.use("Agg")
        plt.ioff()
        torch.cuda.empty_cache()

        num_classes = 2
        model1 = models.resnet18(pretrained=False)
        model2 = models.resnet50(pretrained=False)

        model1.fc = nn.Linear(model1.fc.in_features, num_classes)
        model2.fc = nn.Linear(model2.fc.in_features, num_classes)

        class Ensemble(nn.Module):
            def __init__(self, model1, model2):
                super().__init__()
                self.model1 = model1
                self.model2 = model2
                self.fc = nn.Linear(2 * num_classes, num_classes)

            def forward(self, x):
                if isinstance(x, (list, tuple)):
                    x1 = x[0]
                    x2 = x[1] if len(x) > 1 else x[0]
                else:
                    x1 = x
                    x2 = x
                y1 = self.model1(x1)
                y2 = self.model2(x2)
                y = torch.cat((y1, y2), dim=1)
                return self.fc(y)

        model = Ensemble(model1, model2)

        device = "cuda" if torch.cuda.is_available() else "cpu"
        if device != "cpu":
            model = model.to(device)
        model.load_state_dict(torch.load(file_model, map_location=device))
        model.eval()

        cv2.destroyAllWindows()
        gc.collect()

        print("Инициализация модели завершена")
        print()
        return model

    except Exception as exc:
        print(str(exc))
        return None


def inference_func_ensemble(data=None, model=None, mapping=None, shablon=""):
    try:
        cv2.destroyAllWindows()
        gc.collect()
        torch.cuda.empty_cache()

        device = "cuda" if torch.cuda.is_available() else "cpu"
        if isinstance(data, (list, tuple)) and len(data) >= 2:
            inputs = [
                torch.unsqueeze(torch.tensor(data[0]).cpu(), 0).to(device).float(),
                torch.unsqueeze(torch.tensor(data[1]).cpu(), 0).to(device).float(),
            ]
        else:
            tensor = torch.unsqueeze(torch.tensor(data).cpu(), 0).to(device).float()
            inputs = [tensor, tensor]

        with torch.no_grad():
            output = model(inputs)
            _, predict = torch.max(output.data, 1)

        prediction = mapping[int(np.asarray(predict.cpu())[0])]
        print("PREDICTION" + shablon + ": " + prediction)

        output = output.cpu()
        label = np.asarray(np.argmax(output, axis=1))[0]
        output = np.asarray(torch.squeeze(output, 0))
        expon = np.exp(output - np.max(output))
        probability = round((expon / expon.sum())[label], 2)

        cv2.destroyAllWindows()
        gc.collect()
        print("Уверенность" + shablon + " в предсказании: " + str(probability))
        print("Инференс завершен")
        print()
        return [prediction, probability]

    except Exception as exc:
        print(str(exc))
        return None


def post_func_ensemble(src="", model_type="", prediction="", model_id=0, ind_inference=0, data=None):
    try:
        import matplotlib.pyplot as plt

        matplotlib.use("Agg")
        plt.ioff()

        if int(ind_inference) <= 100 and isinstance(data, (list, tuple)) and len(data) >= 2:
            fig, ax = plt.subplots()
            ax.imshow(np.moveaxis(data[0], 0, -1))
            plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_real_" + str(model_id) + "_" + model_type + ".png")
            plt.clf()
            plt.cla()
            plt.close(fig)
            cv2.destroyAllWindows()
            gc.collect()

            fig, ax = plt.subplots()
            ax.imshow(np.moveaxis(data[1], 0, -1))
            plt.savefig(src + "_inference_" + str(ind_inference) + "_" + prediction + "_mod_" + str(model_id) + "_" + model_type + ".png")
            plt.clf()
            plt.cla()
            plt.close(fig)
            cv2.destroyAllWindows()
            gc.collect()

        plt.clf()
        plt.cla()
        plt.close()
        cv2.destroyAllWindows()
        gc.collect()

        print("Постобработка завершена")
        print()

    except Exception as exc:
        print(str(exc))
        return None