from arr_gtr.grasp import GraspTypeRecognitionModule
import numpy as np
from PIL import Image
from arr_gtr.grasp import grasp_types

nn = GraspTypeRecognitionModule()
#nn = GraspTypeRecognitionModule(pretrained_model = <path to a pth model>)

img = np.asarray(Image.open('sample/image.jpg'))
result = nn.inference(img)
print('Inference result (without affordance): ' + grasp_types[np.argmax(result)])

result = nn.inference_with_affordance(img, 'Apple', affordance_type='varied')
print('Inference result (inference with varied affordance): ' + grasp_types[np.argmax(result)])

result = nn.inference_with_affordance(img, 'Apple', affordance_type='uniformal')
print('Inference result (inference with uniformal affordance): ' + grasp_types[np.argmax(result)])
        
result = nn.inference_from_affordance('Apple', affordance_type='varied')
print('Inference result (inference only with varied affordance): ' + grasp_types[np.argmax(result)])

result = nn.inference_from_affordance('Apple', affordance_type='uniformal')
print('Inference result (inference only with uniformal affordance): ' + grasp_types[np.argmax(result)])

from arr_gtr.grasp import ModelTrainingModule
from arr_gtr.grasp import grasp_types

nn = ModelTrainingModule(<path to a dataset directory>, pretrained_model_path=None, batch_size=512)
nn.train(num_epochs=200, save_path=<path to a pth model e.g., './out/model.pth'>)

nn.visualize_result()