File size: 1,122 Bytes
b6ad7e1 f6ec7b1 b6ad7e1 | 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 29 30 31 32 33 34 35 36 37 | import ultralytics
ultralytics.checks()
from render_results import annotate_image_prediction
from PIL import Image
if __name__=='__main__':
# load model
model = ultralytics.YOLO('../model_v3_202402021.pt')
# set parameters
image_path = '../images/1688033955437.jpg'
saving_folder = '../images'
# infer
image = Image.open(image_path)
results = model.predict(image_path)
# plot
hex_class_colors = {'green_cherry': '#9CF09A',
'yellow_cherry': '#F3C63D',
'red_cherry': '#F44336',
'dark_brown_cherry': '#C36105',
'low_visibility_unsure': '#02D5FA'}
class_dic = {1: 'green_cherry',
3: 'yellow_cherry',
2: 'red_cherry',
0: 'dark_brown_cherry'}
render = annotate_image_prediction(
image_path=image_path,
yolo_boxes=results[0].boxes,
class_dic=class_dic,
saving_folder=saving_folder,
hex_class_colors=hex_class_colors,
show=True,
font_scale=5,
font_thickness=10,
) |