raab/efishSignalDetector
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

bf

Browse Source
This commit is contained in:
Till Raab 2023-10-26 08:47:30 +02:00
parent 68b3967a21
commit 47cd574354
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
inference.py
View File

@ -20,8 +20,8 @@ def plot_inference(img_tensor, output, target, detection_threshold):
quit() quit()
fig, ax = plt.subplots(figsize=(7, 7), num=target['image_id']) fig, ax = plt.subplots(figsize=(7, 7), num=target['image_id'])
ax.imshow(img_tensor.squeeze().permute(1, 2, 0), aspect='auto') ax.imshow(img_tensor.cpu().squeeze().permute(1, 2, 0), aspect='auto')
for (x0, y0, x1, y1), l, score in zip(output[0]['boxes'].cpu(), output[0]['labels'].cpu(), output[0]['scores'].cpu()): for (x0, y0, x1, y1), l, score in zip(output['boxes'].cpu(), output['labels'].cpu(), output['scores'].cpu()):
if score < detection_threshold: if score < detection_threshold:
continue continue
# print(x0, y0, x1, y1, l) # print(x0, y0, x1, y1, l)
@ -50,7 +50,7 @@ def infere_model(test_loader, model, detection_th=0.8):
prog_bar = tqdm(test_loader, total=len(test_loader)) prog_bar = tqdm(test_loader, total=len(test_loader))
for samples, targets in prog_bar: for samples, targets in prog_bar:
images = list(image.to(DEVICE) for image in samples) images = list(image.to(DEVICE) for image in samples)
targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets] targets = [{k: v for k, v in t.items()} for t in targets]
with torch.inference_mode(): with torch.inference_mode():
outputs = model(images) outputs = model(images)