ImageStore Filename Fix

This commit is contained in:
cafeai 2022-11-06 00:32:15 +09:00
parent 305efa20f2
commit b0b1337bc5
1 changed files with 2 additions and 1 deletions

View File

@ -20,6 +20,7 @@ import time
import itertools import itertools
import numpy as np import numpy as np
import json import json
import re
try: try:
pynvml.nvmlInit() pynvml.nvmlInit()
@ -171,7 +172,7 @@ class ImageStore:
# gets caption by removing the extension from the filename and replacing it with .txt # gets caption by removing the extension from the filename and replacing it with .txt
def get_caption(self, index: int) -> str: def get_caption(self, index: int) -> str:
filename = self.image_files[index].split('.')[0] + '.txt' filename = re.sub('\.[^/.]+$', '', self.image_files[index]) + '.txt'
with open(filename, 'r') as f: with open(filename, 'r') as f:
return f.read() return f.read()