We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
1 parent e0ee476 commit d00e65eCopy full SHA for d00e65e
1 file changed
predict.py
@@ -121,7 +121,8 @@ def main():
121
with torch.no_grad():
122
model = model.eval()
123
for img_path in tqdm(image_files):
124
- img_name = os.path.basename(img_path).split('.')[0]
+ ext = os.path.basename(img_path).split('.')[-1]
125
+ img_name = os.path.basename(img_path)[:-len(ext)-1]
126
img = Image.open(img_path).convert('RGB')
127
img = transform(img).unsqueeze(0) # To tensor of NCHW
128
img = img.to(device)
0 commit comments