def crop(pose_img, person_dets):
# cls_dets : x1, y1, x2, y2, score
cls_dets = np.zeros((1, 4), dtype=np.float32)
# test_data : x, y, w, h, score
test_data = np.zeros((1, 4), dtype=np.float32)
test_data[:] = person_dets[:]
bbox = np.asarray(test_data[0])
cls_dets[0, :4] = np.array(
[bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]])
test_imgs = []
details = []
# cropping
test_img, detail = Preprocessing(pose_img, test_data[0], stage='test')
details.append(detail)
details = np.asarray(details).astype(np.float32)
feed = test_img
data = [feed.transpose(0, 2, 3, 1).astype(np.float32)]
return data, details
def dataiter(train_data):
ind = 0
while True:
batch_data = []
for i in range(cfg.batch_size // cfg.nr_aug):
ind += 1
if ind > len(train_data): ind %= len(train_data)
data = Preprocessing(train_data[i])
batch_data.append(data)
ret = []
# aggregate
for i in range(len(batch_data[0])):
ret.append(np.asarray(
[batch_data[j][i] for j in range(len(batch_data))]))
yield ret
def test_net(tester, logger, dets, det_range):
# here we assume all boxes are pre-processed.
nms_method = 'nms'
nms_thresh = 1.
min_scores = 1e-10
min_box_size = 0. # 8 ** 2
all_res = []
dump_results = []
start_time = time.time()
img_start = det_range[0]
while img_start < det_range[1]:
img_end = img_start + 1
im_info = dets[img_start]
while img_end < det_range[1] and dets[img_end]['image_id'] == im_info['image_id']:
img_end += 1
test_data = dets[img_start:img_end]
img_start = img_end
iter_avg_cost_time = (time.time() - start_time) / (img_end - det_range[0])
print('ran %.ds >> << left %.ds' % (
iter_avg_cost_time * (img_end - det_range[0]), iter_avg_cost_time * (det_range[1] - img_end)))
all_res.append([])
# get box detections
cls_dets = np.zeros((len(test_data), 5), dtype=np.float32)
for i in range(len(test_data)):
bbox = np.asarray(test_data[i]['bbox'])
cls_dets[i, :4] = np.array([bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]])
cls_dets[i, 4] = np.array(test_data[i]['score'])
# nms and filter
keep = np.where((cls_dets[:, 4] >= min_scores) &
((cls_dets[:, 3] - cls_dets[:, 1]) * (cls_dets[:, 2] - cls_dets[:, 0]) >= min_box_size))[0]
cls_dets = cls_dets[keep]
if len(cls_dets) > 0:
if nms_method == 'nms':
keep = gpu_nms(cls_dets, nms_thresh)
elif nms_method == 'soft':
keep = cpu_soft_nms(np.ascontiguousarray(cls_dets, dtype=np.float32), method=2)
else:
assert False
cls_dets = cls_dets[keep]
test_data = np.asarray(test_data)[keep]
if len(keep) == 0:
continue
# crop and detect keypoints
cls_skeleton = np.zeros((len(test_data), cfg.nr_skeleton, 3))
crops = np.zeros((len(test_data), 4))
cfg.batch_size = 32
batch_size = cfg.batch_size // 2
for test_id in range(0, len(test_data), batch_size):
start_id = test_id
end_id = min(len(test_data), test_id + batch_size)
test_imgs = []
details = []
for i in range(start_id, end_id):
test_img, detail = Preprocessing(test_data[i], stage='test')
test_imgs.append(test_img)
details.append(detail)
details = np.asarray(details)
feed = test_imgs
for i in range(end_id - start_id):
ori_img = test_imgs[i][0].transpose(1, 2, 0)
flip_img = cv2.flip(ori_img, 1)
feed.append(flip_img.transpose(2, 0, 1)[np.newaxis, ...])
feed = np.vstack(feed)
res = tester.predict_one([feed.transpose(0, 2, 3, 1).astype(np.float32)])[0]
res = res.transpose(0, 3, 1, 2)
for i in range(end_id - start_id):
fmp = res[end_id - start_id + i].transpose((1, 2, 0))
fmp = cv2.flip(fmp, 1)
fmp = list(fmp.transpose((2, 0, 1)))
for (q, w) in cfg.symmetry:
fmp[q], fmp[w] = fmp[w], fmp[q]
fmp = np.array(fmp)
res[i] += fmp
res[i] /= 2
for test_image_id in range(start_id, end_id):
r0 = res[test_image_id - start_id].copy()
r0 /= 255.
r0 += 0.5
for w in range(cfg.nr_skeleton):
res[test_image_id - start_id, w] /= np.amax(res[test_image_id - start_id, w])
border = 10
dr = np.zeros((cfg.nr_skeleton, cfg.output_shape[0] + 2 * border, cfg.output_shape[1] + 2 * border))
dr[:, border:-border, border:-border] = res[test_image_id - start_id][:cfg.nr_skeleton].copy()
for w in range(cfg.nr_skeleton):
dr[w] = cv2.GaussianBlur(dr[w], (21, 21), 0)
for w in range(cfg.nr_skeleton):
lb = dr[w].argmax()
y, x = np.unravel_index(lb, dr[w].shape)
dr[w, y, x] = 0
lb = dr[w].argmax()
py, px = np.unravel_index(lb, dr[w].shape)
y -= border
x -= border
py -= border + y
px -= border + x
ln = (px ** 2 + py ** 2) ** 0.5
delta = 0.25
if ln > 1e-3:
x += delta * px / ln
y += delta * py / ln
x = max(0, min(x, cfg.output_shape[1] - 1))
y = max(0, min(y, cfg.output_shape[0] - 1))
cls_skeleton[test_image_id, w, :2] = (x * 4 + 2, y * 4 + 2)
cls_skeleton[test_image_id, w, 2] = r0[w, int(round(y) + 1e-10), int(round(x) + 1e-10)]
# map back to original images
crops[test_image_id, :] = details[test_image_id - start_id, :]
for w in range(cfg.nr_skeleton):
cls_skeleton[test_image_id, w, 0] = cls_skeleton[test_image_id, w, 0] / cfg.data_shape[1] * (
crops[test_image_id][2] - crops[test_image_id][0]) + crops[test_image_id][0]
cls_skeleton[test_image_id, w, 1] = cls_skeleton[test_image_id, w, 1] / cfg.data_shape[0] * (
crops[test_image_id][3] - crops[test_image_id][1]) + crops[test_image_id][1]
all_res[-1] = [cls_skeleton.copy(), cls_dets.copy()]
cls_partsco = cls_skeleton[:, :, 2].copy().reshape(-1, cfg.nr_skeleton)
cls_skeleton[:, :, 2] = 1
cls_scores = cls_dets[:, -1].copy()
# rescore
cls_dets[:, -1] = cls_scores * cls_partsco.mean(axis=1)
cls_skeleton = np.concatenate(
[cls_skeleton.reshape(-1, cfg.nr_skeleton * 3), (cls_scores * cls_partsco.mean(axis=1))[:, np.newaxis]],
axis=1)
for i in range(len(cls_skeleton)):
result = dict(image_id=im_info['image_id'], category_id=1, score=float(round(cls_skeleton[i][-1], 4)),
keypoints=cls_skeleton[i][:-1].round(3).tolist())
dump_results.append(result)
return all_res, dump_results
def test_net(tester, logger, dets, det_range):
# here we assume all boxes are pre-processed.
all_res = []
dump_results = []
start_time = time.time()
img_start = det_range[0]
while img_start < det_range[1]:
img_end = img_start + 1
im_info = dets[img_start]
while img_end < det_range[1] and dets[img_end]['imgpath'] == im_info[
'imgpath']:
img_end += 1
test_data = dets[img_start:img_end]
img_start = img_end
iter_avg_cost_time = (time.time() - start_time) / (img_end -
det_range[0])
print('ran %.ds >> << left %.ds' %
(iter_avg_cost_time *
(img_end - det_range[0]), iter_avg_cost_time *
(det_range[1] - img_end)))
all_res.append([])
#test_data = np.asarray(test_data)
# crop and detect keypoints
cls_skeleton = np.zeros((len(test_data), cfg.nr_skeleton, 3))
crops = np.zeros((len(test_data), 4))
cfg.batch_size = 8
batch_size = cfg.batch_size // 2
for test_id in range(0, len(test_data), batch_size):
start_id = test_id
end_id = min(len(test_data), test_id + batch_size)
test_imgs = []
details = []
for i in range(start_id, end_id):
test_img, detail = Preprocessing(test_data[i], stage='test')
test_imgs.append(test_img)
details.append(detail)
details = np.asarray(details)
feed = test_imgs
for i in range(end_id - start_id):
ori_img = test_imgs[i][0].transpose(1, 2, 0)
flip_img = cv2.flip(ori_img, 1)
feed.append(flip_img.transpose(2, 0, 1)[np.newaxis, ...])
feed = np.vstack(feed)
res = tester.predict_one(
[feed.transpose(0, 2, 3, 1).astype(np.float32)])[0]
res = res[0]
res = res.transpose(0, 3, 1, 2)
for i in range(end_id - start_id):
fmp = res[end_id - start_id + i].transpose((1, 2, 0))
fmp = cv2.flip(fmp, 1)
fmp = list(fmp.transpose((2, 0, 1)))
for (q, w) in cfg.symmetry:
fmp[q], fmp[w] = fmp[w], fmp[q]
fmp = np.array(fmp)
res[i] += fmp
res[i] /= 2
for test_image_id in range(start_id, end_id):
r0 = res[test_image_id - start_id].copy()
r0 /= 255.
r0 += 0.5
for w in range(cfg.nr_skeleton):
res[test_image_id - start_id,
w] /= np.amax(res[test_image_id - start_id, w])
border = 10
dr = np.zeros(
(cfg.nr_skeleton, cfg.output_shape[0] + 2 * border,
cfg.output_shape[1] + 2 * border))
dr[:, border:-border,
border:-border] = res[test_image_id -
start_id][:cfg.nr_skeleton].copy()
for w in range(cfg.nr_skeleton):
dr[w] = cv2.GaussianBlur(dr[w], (21, 21), 0)
for w in range(cfg.nr_skeleton):
lb = dr[w].argmax()
y, x = np.unravel_index(lb, dr[w].shape)
dr[w, y, x] = 0
lb = dr[w].argmax()
py, px = np.unravel_index(lb, dr[w].shape)
y -= border
x -= border
py -= border + y
px -= border + x
ln = (px**2 + py**2)**0.5
delta = 0.25
if ln > 1e-3:
x += delta * px / ln
y += delta * py / ln
x = max(0, min(x, cfg.output_shape[1] - 1))
y = max(0, min(y, cfg.output_shape[0] - 1))
cls_skeleton[test_image_id, w, :2] = (x * 4 + 2, y * 4 + 2)
cls_skeleton[test_image_id, w,
2] = r0[w,
int(round(y) + 1e-10),
int(round(x) + 1e-10)]
# map back to original images
crops[test_image_id, :] = details[test_image_id - start_id, :]
for w in range(cfg.nr_skeleton):
cls_skeleton[test_image_id, w, 0] = cls_skeleton[
test_image_id, w, 0] / cfg.data_shape[1] * (
crops[test_image_id][2] -
crops[test_image_id][0]) + crops[test_image_id][0]
cls_skeleton[test_image_id, w, 1] = cls_skeleton[
test_image_id, w, 1] / cfg.data_shape[0] * (
crops[test_image_id][3] -
crops[test_image_id][1]) + crops[test_image_id][1]
all_res[-1] = cls_skeleton.copy()
#cls_partsco = cls_skeleton[:, :, 2].copy().reshape(-1, cfg.nr_skeleton)
#cls_skeleton[:, :, 2] = 1
# rescore
cls_skeleton = np.concatenate(
[cls_skeleton.reshape(-1, cfg.nr_skeleton * 3)], axis=1)
for i in range(len(cls_skeleton)):
result = dict(keypoints=cls_skeleton[i][:].round().tolist())
dump_results.append(result)
return all_res, dump_results
def main():
axis = 'ax1'
# CUDA for PyTorch
device = train_device()
# Training dataset
train_params = {'batch_size': 10, 'shuffle': True, 'num_workers': 4}
data_path = './dataset/dataset_' + axis + '/train/'
train_dataset = Dataset(data_path,
transform=transforms.Compose([Preprocessing()]))
lenght = int(len(train_dataset))
train_loader = torch.utils.data.DataLoader(train_dataset, **train_params)
# Validation dataset
data_path = './dataset/dataset_' + axis + '/valid/'
valid_dataset = Dataset(data_path,
transform=transforms.Compose([Preprocessing()]))
valid_params = {'batch_size': 10, 'shuffle': True, 'num_workers': 4}
val_loader = torch.utils.data.DataLoader(valid_dataset, **valid_params)
# Training params
learning_rate = 1e-4
max_epochs = 100
# Used pretrained model and modify channels from 3 to 1
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch',
'unet',
in_channels=3,
out_channels=1,
init_features=32,
pretrained=True)
model.encoder1.enc1conv1 = nn.Conv2d(1,
32,
kernel_size=(3, 3),
stride=(1, 1),
padding=(1, 1),
bias=False)
model.to(device)
# Optimizer and loss function
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
dsc_loss = DiceLoss()
# Metrics
train_loss = AverageMeter('Training loss', ':.6f')
val_loss = AverageMeter('Validation loss', ':.6f')
best_loss = float('inf')
nb_of_batches = lenght // train_params['batch_size']
for epoch in range(max_epochs):
val_loss.avg = 0
train_loss.avg = 0
if not epoch:
logg_file = loggs.Loggs(['epoch', 'train_loss', 'val_loss'])
model.train()
for i, (image, label) in enumerate(train_loader):
torch.cuda.empty_cache()
image, label = image.to(device), label.to(device)
optimizer.zero_grad()
y_pred = model(image)
loss = dsc_loss(y_pred, label)
del y_pred
train_loss.update(loss.item(), image.size(0))
loss.backward()
optimizer.step()
loggs.training_bar(i,
nb_of_batches,
prefix='Epoch: %d/%d' % (epoch, max_epochs),
suffix='Loss: %.6f' % loss.item())
print(train_loss.avg)
with torch.no_grad():
for i, (x_val, y_val) in enumerate(val_loader):
x_val, y_val = x_val.to(device), y_val.to(device)
model.eval()
yhat = model(x_val)
loss = dsc_loss(yhat, y_val)
val_loss.update(loss.item(), x_val.size(0))
print(val_loss)
logg_file.save([epoch, train_loss.avg, val_loss.avg])
# Save the best model with minimum validation loss
if best_loss > val_loss.avg:
print('Updated model with validation loss %.6f ---> %.6f' %
(best_loss, val_loss.avg))
best_loss = val_loss.avg
torch.save(model, './model_' + axis + '/best_model.pt')
def _test_net(self, dets):
# here we assume all boxes are pre-processed.
det_range = [0, len(dets)]
nms_method = 'nms'
nms_thresh = 1.
min_scores = 0.5 # 1e-10 modified to avoid mismatch
min_box_size = 0. # 8 ** 2
all_res = []
dump_results = []
start_time = time.time()
img_start = det_range[0]
while img_start < det_range[1]:
img_end = img_start + 1
im_info = dets[img_start]
# toate img_id au acelasi id.
#din cauza asta, img_end o sa fie egal cu numarul de detectii
while img_end < det_range[1] and dets[img_end][
'image_id'] == im_info['image_id']:
img_end += 1
#practic tot
test_data = dets[img_start:img_end]
img_start = img_end
all_res.append([])
# get box detections
#face o matrice de zero-uri cu 5 coloane si linii cate elemente sunt
cls_dets = np.zeros((len(test_data), 5), dtype=np.float32)
for i in range(len(test_data)):
bbox = np.asarray(test_data[i]['bbox'])
#banuiesc ca asta e formula cu care salveaza bbox, e cam dubioasa
cls_dets[i, :4] = np.array(
[bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]])
#pe ultima pozitie e scorul de incredere
cls_dets[i, 4] = np.array(test_data[i]['score'])
# nms and filter
keep = np.where((cls_dets[:, 4] >= min_scores) & (
(cls_dets[:, 3] - cls_dets[:, 1]) *
(cls_dets[:, 2] - cls_dets[:, 0]) >= min_box_size))[0]
#se pastreaza doar alea peste treshold
cls_dets = cls_dets[keep]
if len(cls_dets) > 0:
if nms_method == 'nms':
keep = gpu_nms(cls_dets, nms_thresh)
elif nms_method == 'soft':
keep = cpu_soft_nms(np.ascontiguousarray(cls_dets,
dtype=np.float32),
method=2)
else:
assert False
cls_dets = cls_dets[keep]
test_data = np.asarray(test_data)[keep]
if len(keep) == 0:
continue
# crop and detect keypoints
#cls_skeleton o sa fie un cub
cls_skeleton = np.zeros((len(test_data), cfg.nr_skeleton, 3))
#crops o sa fie o matrice cu len linii si 4 coloane
crops = np.zeros((len(test_data), 4))
cfg.batch_size = 32
#practic batch-sizeul o sa fie 16. faza e ca as putea sa fac asta foarte mic, dar nu va afecta cu nimic, ca problema nu e aici
batch_size = cfg.batch_size // 2
#de la 0 la len, din 16 in 16
for test_id in range(0, len(test_data), batch_size):
start_id = test_id
#end-ul va fi mereu inceput+batch size
end_id = min(len(test_data), test_id + batch_size)
test_imgs = []
details = []
#spre exemplu e de la 0 la 16, sau de la 1 la 17
for i in range(start_id, end_id):
#merg in preprocessing sa vad ke se intampla. se trimite o imagine si stageul de test
#am venit din Preprocessing. Se primeste iamginea normalizata si cu ceva border, si in detail sunt dimensiuni
test_img, detail = Preprocessing(test_data[i],
stage='test')
#in test_imgs vor fi imaginile cropuite ale persoanelor din bounding boxes
test_imgs.append(test_img)
details.append(detail)
details = np.asarray(details)
feed = test_imgs
for i in range(end_id - start_id):
#se transpun imaginile, nu stiu cum
ori_img = test_imgs[i][0].transpose(1, 2, 0)
flip_img = cv2.flip(ori_img, 1)
feed.append(flip_img.transpose(2, 0, 1)[np.newaxis, ...])
feed = np.vstack(feed)
#in feed ar trebui sa fie imaginile procesate si imaginile flipuite si transpuse. e un array cu imaginile astea din batchul actual
#merg sa vad ce face predict_one. primeste imaginile transpuse inca o data din batchul actual procesate
res = self.tester.predict_one(
[feed.transpose(0, 2, 3, 1).astype(np.float32)])[0]
# from IPython import embed; embed()
res = res.transpose(0, 3, 1, 2)
#practic nu imaginile specifice, ci de la 0 la diferenta, care probabil e 16 mereu, adica range(16)
for i in range(end_id - start_id):
#ia o imagine de genul 16+i, care i e de la 0 la 16 si o transpune din nou.
fmp = res[end_id - start_id + i].transpose((1, 2, 0))
fmp = cv2.flip(fmp, 1)
fmp = list(fmp.transpose((2, 0, 1)))
#inverseaza pixeli poate
for (q, w) in cfg.symmetry:
fmp[q], fmp[w] = fmp[w], fmp[q]
fmp = np.array(fmp)
#aici chiar nu mai inteleg ce face, de ce ia o imagine din primul batch si aduna la ea imaginea fmp care e mereu din al doilea batch
res[i] += fmp
res[i] /= 2
heatmaps = []
for test_image_id in range(start_id, end_id):
#din nou, r0 va fi mereu din primele 16
r0 = res[test_image_id - start_id].copy()
#o normalizeaza cu std si mean
r0 /= 255.
r0 += 0.5
#cfg.nr_skeleton e 17
for w in range(cfg.nr_skeleton):
#din aceasta linie trag concluzia ca res e un cub. si imparte acest array cu maximul din el
res[test_image_id - start_id,
w] /= np.amax(res[test_image_id - start_id, w])
border = 10
#in dr pune toate pozele din primul batch
dr = np.zeros(
(cfg.nr_skeleton, cfg.output_shape[0] + 2 * border,
cfg.output_shape[1] + 2 * border))
dr[:, border:-border, border:-border] = res[
test_image_id - start_id][:cfg.nr_skeleton].copy()
# TODO: try to use those with out gaussian
for w in range(cfg.nr_skeleton):
#blureaza imaginile
dr[w] = cv2.GaussianBlur(dr[w], (21, 21), 0)
# dr[w] = cv2.GaussianBlur ( dr[w], (1, 1), 0 ) # Will working on it.
raw_heatmaps = list(dr[:, border:-border,
border:-border].copy())
#de la 0 la 17
for w in range(cfg.nr_skeleton):
#indicele celei mai mari valori din dr[w]
lb = dr[w].argmax()
#nu prea inteleg aceasta functie, ideea e ca transforma indicii din lb in 2 array-uri de coordonate
y, x = np.unravel_index(lb, dr[w].shape)
#pune pe 0 astea ce au fost deja procesate
dr[w, y, x] = 0
#mai jos sunt niste operatii matematice pe care nu o sa le pricep acum, sunt doar calcule pentru a obtine coordonatele corecte de pe heatmaps
lb = dr[w].argmax()
py, px = np.unravel_index(lb, dr[w].shape)
y -= border
x -= border
py -= border + y
px -= border + x
ln = (px**2 + py**2)**0.5
delta = 0.25
if ln > 1e-3:
x += delta * px / ln
y += delta * py / ln
x = max(0, min(x, cfg.output_shape[1] - 1))
y = max(0, min(y, cfg.output_shape[0] - 1))
cls_skeleton[test_image_id,
w, :2] = (x * 4 + 2, y * 4 + 2)
cls_skeleton[test_image_id, w,
2] = r0[w,
int(round(y) + 1e-10),
int(round(x) + 1e-10)]
# map back to original images
#practic in crops sunt dimensiunile originale ale imaginilor care erau si in details
crops[test_image_id, :] = details[test_image_id -
start_id, :]
length = crops[test_image_id][2] - crops[test_image_id][0]
width = crops[test_image_id][3] - crops[test_image_id][1]
#cred ca e orientarea salvata in astea 2
l_ori = raw_heatmaps[0].shape[0]
w_ori = raw_heatmaps[0].shape[1]
def test_for_loop(raw_heatmaps, order=3):
heatmaps_this = copy.deepcopy(raw_heatmaps)
for w in range(cfg.nr_skeleton):
heatmaps_this[
w] = scipy.ndimage.interpolation.zoom(
raw_heatmaps[w],
(width / l_ori, length / w_ori),
order=order)
heatmaps_this[w][heatmaps_this[w] <= 0] = 10e-6
zoomed_heatmaps = np.array(heatmaps_this)
return zoomed_heatmaps
def test_my(raw_heatmaps, order=3):
origin_heatmap = np.array(raw_heatmaps)
zoomed_heatmaps = scipy.ndimage.interpolation.zoom(
origin_heatmap, (1, width / l_ori, length / w_ori),
order=order)
zoomed_heatmaps[zoomed_heatmaps <= 0] = 10e-6
return zoomed_heatmaps
# origin_heatmap = np.array ( raw_heatmaps )
zoomed_heatmaps = np.empty(
(len(raw_heatmaps), int(width), int(length)))
for zoom_id, heatmap_i in enumerate(raw_heatmaps):
zoomed_heatmaps[
zoom_id] = scipy.ndimage.interpolation.zoom(
heatmap_i, (width / l_ori, length / w_ori),
order=3)
# zoomed_heatmaps = np.array ( raw_heatmaps )
zoomed_heatmaps[zoomed_heatmaps <= 0] = 10e-6
# orinial_heatmaps = np.zeros((cfg.nr_skeleton, ori_width, ori_length))
# orinial_heatmaps[:, int(crops[test_image_id][1]):int(crops[test_image_id][3]), int(crops[test_image_id][0]):int(crops[test_image_id][2])] = heatmaps_this
# orinial_heatmaps[orinial_heatmaps <= 0] = 10**-6
heatmaps.append(zoomed_heatmaps)
for w in range(cfg.nr_skeleton):
cls_skeleton[test_image_id, w, 0] = cls_skeleton[
test_image_id, w, 0] / cfg.data_shape[1] * (
crops[test_image_id][2] - crops[test_image_id]
[0]) + crops[test_image_id][0]
cls_skeleton[test_image_id, w, 1] = cls_skeleton[
test_image_id, w, 1] / cfg.data_shape[0] * (
crops[test_image_id][3] - crops[test_image_id]
[1]) + crops[test_image_id][1]
all_res[-1] = [cls_skeleton.copy(), cls_dets.copy()]
cls_partsco = cls_skeleton[:, :,
2].copy().reshape(-1, cfg.nr_skeleton)
cls_skeleton[:, :, 2] = 1
cls_scores = cls_dets[:, -1].copy()
# rescore
cls_dets[:, -1] = cls_scores * cls_partsco.mean(axis=1)
cls_skeleton = np.concatenate([
cls_skeleton.reshape(-1, cfg.nr_skeleton * 3),
(cls_scores * cls_partsco.mean(axis=1))[:, np.newaxis]
],
axis=1)
for i in range(len(cls_skeleton)):
result = dict(image_id=im_info['image_id'],
category_id=1,
score=float(round(cls_skeleton[i][-1], 4)),
keypoints=cls_skeleton[i][:-1].round(3).tolist(),
bbox=dets[i]['bbox'],
heatmaps=heatmaps[i],
crops=crops[i])
dump_results.append(result)
if self.show_image:
import pdb
pdb.set_trace()
dbg_im = dets[0][
'data'] # Since all detection are based on the same image
from utils.visualize import visualize
visualize(dbg_im,
keypoints=[i['keypoints'] for i in dump_results],
det_boxes=cls_dets)
# import pdb; pdb.set_trace ()
# return all_res, dump_results
return dump_results
def main():
device = train_device()
# Test dataset
test_params = {'batch_size': 5, 'shuffle': False, 'num_workers': 4}
data_path = './dataset/test/'
kod = [f for f in os.listdir(data_path)]
samples = [data_path + f + '/ax1/' for f in os.listdir(data_path)]
info_file = [data_path + f + '/info.txt' for f in os.listdir(data_path)]
# Load model
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch',
'unet',
in_channels=3,
out_channels=1,
init_features=32,
pretrained=True)
model.encoder1.enc1conv1 = nn.Conv2d(1,
32,
kernel_size=(3, 3),
stride=(1, 1),
padding=(1, 1),
bias=False)
model = torch.load('./model_ax1/best_model.pt', map_location=device)
model.to(device)
model.eval()
j = 0
# Testing
for sample in samples:
with open(info_file[j], 'r') as f:
info = f.read()
info = info.replace(",", "")
img_shape = list(map(int, (info.split()[-3:])))
test_dataset = Dataset(sample,
is_test=True,
transform=transforms.Compose([Preprocessing()]))
test_loader = torch.utils.data.DataLoader(test_dataset, **test_params)
result = np.zeros(tuple(img_shape))
affine = return_affine(kod[j])
k = 0
for i, (x_test) in enumerate(test_loader):
torch.cuda.empty_cache()
x_test = x_test.to(device)
outputs = model(x_test)
l = 0
for out in outputs:
key = ord('a')
img = np.array(np.squeeze((out.cpu().detach().numpy())),
np.uint8)
img2 = np.array(
np.squeeze(255 * (x_test[l].cpu().detach().numpy())),
np.uint8)
img = crop(img, (img_shape[0], img_shape[1]))
img2 = crop(img2, (img_shape[0], img_shape[1]))
print(max(img.flatten()))
img = np.flip(img, axis=0)
img = np.array(np.where(img < 0.5, 0, 1), np.float32)
print(max(img.flatten()))
result[:, :, k] = img
l += 1
k += 1
save_labels(result, affine,
Path('./predictions2_ax1/' + kod[j] + '.nii.gz'))
j += 1
def main():
# CUDA for PyTorch
use_cuda = torch.cuda.is_available()
device = torch.device('cuda:0' if use_cuda else 'cpu')
torch.backends.cudnn.benchmark = True
train_params = {'batch_size': 50, 'shuffle': True, 'num_workers': 4}
valid_params = {'batch_size': 100, 'shuffle': True, 'num_workers': 4}
# Load dataset
data_path = '../generated_data/'
my_dataset = Dataset(data_path,
transform=transforms.Compose([Preprocessing()]))
lengths = [int(len(my_dataset) * 0.8), int(len(my_dataset) * 0.2)]
train_dataset, val_dataset = random_split(my_dataset, lengths)
train_loader = torch.utils.data.DataLoader(train_dataset, **train_params)
val_loader = torch.utils.data.DataLoader(val_dataset, **valid_params)
# Training params
learning_rate = 1e-3
max_epochs = 4
# Model
model = unet.ResUNet(2, 1, n_size=16)
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
train_loss = AverageMeter('Training loss', ':.6f')
val_loss = AverageMeter('Validation loss', ':.6f')
best_loss = float('inf')
nb_of_batches = lengths[0] // train_params['batch_size']
# Training loop
for epoch in range(max_epochs):
if not epoch:
logg_file = loggs.Loggs(['epoch', 'train_loss', 'val_loss'])
for i, (x_batch, y_labels) in enumerate(train_loader):
x_batch, y_labels = x_batch.to(device), y_labels.to(device)
y_pred = model(x_batch)
#y_pred = torch.round(y_pred[0])
loss = dice_loss(y_pred, y_labels)
train_loss.update(loss.item(), x_batch.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
loggs.training_bar(i,
nb_of_batches,
prefix='Epoch: %d/%d' % (epoch, max_epochs),
suffix='Loss: %.6f' % loss.item())
print(train_loss)
with torch.no_grad():
for i, (x_val, y_val) in enumerate(val_loader):
x_val, y_val = x_val.to(device), y_val.to(device)
model.eval()
yhat = model(x_val)
loss = dice_loss(yhat, y_val)
val_loss.update(loss.item(), x_val.size(0))
if i == 10: break
print(val_loss)
logg_file.save([epoch, train_loss.avg, val_loss.avg])
# Save the best model with minimum validation loss
if best_loss > val_loss.avg:
print('Updated model with validation loss %.6f ---> %.6f' %
(best_loss, val_loss.avg))
best_loss = val_loss.avg
torch.save(model, 'best_model.pt')