Views
No views yet
1from huggingface_hub import hf_hub_download
2import cv2
3import numpy as np
4import PIL.Image
5import torch
6import torchvision
7
8DEVICE_NAME = 'cuda:0'
9MODEL_PATH = hf_hub_download('kindwise/router.small', 'model.traced.pt')
10CLASSES_PATH = hf_hub_download('kindwise/router.small', 'classes.txt')
11IMAGE_PATH = '/tmp/photo.jpg'
12
13with open(CLASSES_PATH) as f:
14 CLASSES = [line.strip() for line in f]
15MODEL = torch.jit.load(MODEL_PATH).eval().to(DEVICE_NAME)
16
17
18def resize_crop(image_data: np.ndarray, target_size: int = 480) -> np.ndarray | None:
19 height, width, _ = image_data.shape
20 # Determine the size of the square crop
21 crop_size = min(height, width)
22 # Calculate coordinates for center crop
23 start_x = (width - crop_size) // 2
24 start_y = (height - crop_size) // 2
25 # Perform center crop
26 cropped_img = image_data[
27 start_y : start_y + crop_size,
28 start_x : start_x + crop_size
29 ]
30 # Resize cropped image to target size
31 return cv2.resize(
32 cropped_img,
33 (target_size, target_size),
34 interpolation=cv2.INTER_AREA,
35 )
36
37with torch.no_grad():
38 image_array = np.array(PIL.Image.open(IMAGE_PATH))
39 image_array_resized = resize_crop(image_array)
40 image_tensor = torchvision.transforms.functional.to_tensor(image_array_resized).to(DEVICE_NAME)
41 prediction = MODEL(image_tensor.unsqueeze(0)).squeeze(0).cpu().numpy()
42 for i in (-prediction).argsort():
43 print(f'{CLASSES[i]:>10}: {100 * prediction[i]:.1f}%') plant: 91.3%
unhealthy_plant: 53.3%
crop: 16.2%
insect: 0.4%
human: 0.1%
mushroom: 0.0%1from huggingface_hub import hf_hub_download
2import numpy as np
3import tensorflow as tf
4
5MODEL_PATH = hf_hub_download('kindwise/router.small', 'model.tflite') # or model.optimized.tflite
6CLASSES_PATH = hf_hub_download('kindwise/router.small', 'classes.txt')
7
8with open(CLASSES_PATH) as f:
9 CLASSES = [line.strip() for line in f]
10INTERPRETER = tf.lite.Interpreter(model_path=MODEL_PATH)
11INTERPRETER.allocate_tensors()
12
13image_array_resized = ... # see the previous example
14tf_input = np.expand_dims( # add batch dimension
15 (image_array_resized / 255).astype(np.float32), # image values in [0..1]
16 0,
17)
18input_details = INTERPRETER.get_input_details()
19output_details = INTERPRETER.get_output_details()
20INTERPRETER.set_tensor(
21 input_details[0]['index'],
22 tf_input,
23)
24INTERPRETER.invoke()
25logits = INTERPRETER.get_tensor(output_details[0]['index'])[0]
26prediction = tf.nn.sigmoid(logits).numpy()
27for i in (-prediction).argsort():
28 print(f'{CLASSES[i]:>10}: {100 * prediction[i]:.1f}%') plant: 91.3%
unhealthy_plant: 53.3%
crop: 16.2%
insect: 0.4%
human: 0.1%
mushroom: 0.0%