def test(path): model = Model() model.to("cuda:0") model.eval() checkpoint = torch.load("./model.pth") model.load_state_dict(checkpoint["model"]) img = np.array(Image.open(path).resize([448, 448]))[np.newaxis] img = np.transpose(img, axes=[0, 3, 1, 2]) / 255 img = torch.tensor(img, dtype=torch.float32).to("cuda:0") preds = model(img).cpu().detach().numpy() cell_h, cell_w = IMG_H / S, IMG_W / S x, y = np.meshgrid(range(S), range(S)) preds_xywhs = [] for i in range(B): preds_x = (preds[0, :, :, i * 4] + x) * cell_w preds_y = (preds[0, :, :, i * 4 + 1] + y) * cell_h preds_w = preds[0, :, :, i * 4 + 2] * IMG_W preds_h = preds[0, :, :, i * 4 + 3] * IMG_H preds_xywh = np.dstack((preds_x, preds_y, preds_w, preds_h)) preds_xywhs.append(preds_xywh) preds_xywhs = np.dstack(preds_xywhs) preds_xywhs = np.reshape(preds_xywhs, [-1, 4]) preds_class = preds[0, :, :, 10:] preds_class = np.reshape(preds_class, [-1, 20]) preds_c = preds[0, :, :, 8:10] preds_c = np.reshape(preds_c, [-1, 1]) max_arg = np.argmax(preds_c, axis=0) print("max confidence: %f" % (preds_c[max_arg])) max_arg_ = np.argmax(preds_class[int(max_arg // 2)]) print("class confidence: %f" % (preds_class[max_arg // 2, max_arg_])) print("class category: %s" % (CLASSES[int(max_arg_)])) Image.fromarray( np.uint8( draw_bboxes(np.array(Image.open(path).resize([448, 448])), preds_xywhs[max_arg[0]:max_arg[0] + 1]))).show()
#!/usr/bin/env python3 # -*- coding: utf-8 -*- # добавлем родительский каталог в пути посика модуле # чтобы был виден модуль utils import os import sys sys.path = [os.pardir] + sys.path # чтобы можно было загрузить изображение по частям from utils import Image # чтобы перемешать части изображения import random orig = Image.open('cat.png') img = Image.new(orig.x_size, orig.y_size) parts = [orig[i] for i in range(len(orig))] # parts = [image for image in orig] random.shuffle(parts) for i in range(len(orig)): img[i] = parts[i] img.save('cat_break.png') img.show('cat_break')