Skip to content

Commit bc0cf9f

Browse files
mkhoshbin72m.khoshbin
and
m.khoshbin
authored
some modifications added to rtdetrv2_torch.py (#399)
Co-authored-by: m.khoshbin <[email protected]>
1 parent 2c65ad5 commit bc0cf9f

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

rtdetrv2_pytorch/references/deploy/rtdetrv2_torch.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import numpy as np
99
from PIL import Image, ImageDraw
1010

11-
from cvperception.core import YAMLConfig
11+
from src.core import YAMLConfig
1212

1313

1414
def draw(images, labels, boxes, scores, thrh = 0.6):
@@ -18,10 +18,11 @@ def draw(images, labels, boxes, scores, thrh = 0.6):
1818
scr = scores[i]
1919
lab = labels[i][scr > thrh]
2020
box = boxes[i][scr > thrh]
21+
scrs = scores[i][scr > thrh]
2122

22-
for b in box:
23+
for j,b in enumerate(box):
2324
draw.rectangle(list(b), outline='red',)
24-
draw.text((b[0], b[1]), text=str(lab[i].item()), fill='blue', )
25+
draw.text((b[0], b[1]), text=f"{lab[j].item()} {round(scrs[j].item(),2)}", fill='blue', )
2526

2627
im.save(f'results_{i}.jpg')
2728

0 commit comments

Comments
 (0)