Spaces:
Running
Running
File size: 481 Bytes
dbdb9fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
import torch
from headshot import Headshot
from headshot import config
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = Headshot().to(device)
pretrained = None
if pretrained:
model_path = ''
pass
else:
model_path = ''
pass
model.load_state_dict(model_path)
def sample():
image_path = './interface/images/demo.jpg'
prediction,image = model.predict_image(image_path)
print(f"Prediction ->{prediction}")
return prediction,image |