добавил модели inference на двух картинках
Sergey Revyakin
2 days ago
parent
0fad5d6404
commit
a1c99ebf9f
@ -0,0 +1,197 @@
|
||||
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
|
||||
@ -0,0 +1,197 @@
|
||||
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
|
||||
Loading…
Reference in New Issue