test run prediction on image
This commit is contained in:
+30
-2
@@ -1,7 +1,12 @@
|
|||||||
"""Main script to be run by service."""
|
"""Main script to be run by service."""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from traceback import print_exc
|
||||||
|
|
||||||
|
import torch
|
||||||
from models import VisualCommunicationModel
|
from models import VisualCommunicationModel
|
||||||
from utils import load_model
|
from tqdm import tqdm
|
||||||
|
from utils import DEVICE, VCDADataset, load_model
|
||||||
|
|
||||||
from shared.data_store import connect
|
from shared.data_store import connect
|
||||||
from shared.utils import setup_logging
|
from shared.utils import setup_logging
|
||||||
@@ -13,6 +18,29 @@ if __name__ == '__main__':
|
|||||||
minio_client = connect()
|
minio_client = connect()
|
||||||
# instantiate model
|
# instantiate model
|
||||||
model: VisualCommunicationModel = load_model(client=minio_client)
|
model: VisualCommunicationModel = load_model(client=minio_client)
|
||||||
# get image
|
model.eval()
|
||||||
|
# setup dataset
|
||||||
|
dataset = VCDADataset(
|
||||||
|
data_name_list=[
|
||||||
|
'02dbaf48d713e4e6d3a6b98fd2dc866e',
|
||||||
|
],
|
||||||
|
do_augment=False,
|
||||||
|
)
|
||||||
|
|
||||||
# make prediction
|
# make prediction
|
||||||
|
with torch.no_grad():
|
||||||
|
for i in tqdm(range(len(dataset))):
|
||||||
|
try:
|
||||||
|
# get image
|
||||||
|
image = dataset.__getitem__(i)
|
||||||
|
image = torch.unsqueeze(image, 0) # add artificial batch dimension
|
||||||
|
image = image.to(DEVICE)
|
||||||
|
# make prediction
|
||||||
|
pred = model(image)
|
||||||
|
pred = pred.cpu().numpy()
|
||||||
|
except Exception:
|
||||||
|
print_exc()
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
logging.info(pred)
|
||||||
|
logging.debug('finished')
|
||||||
|
|||||||
Reference in New Issue
Block a user