Weather_Image_Classification-Capstone_Project/lambda_function.py at main · cjj1120/Weather_Image_Classification-Capstone_Project · GitHub

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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import tflite_runtime.interpreter as tflite
from io import BytesIO
from urllib import request
from PIL import Image
import numpy as np


def download_image(url):
    with request.urlopen(url) as resp:
        buffer = resp.read()
    stream = BytesIO(buffer)
    img = Image.open(stream)
    return img


def prepare_image(img):
    if img.mode != 'RGB':
        img = img.convert('RGB')
    img = img.resize((224, 224), Image.NEAREST)
    return img


interpreter =  tflite.Interpreter(model_path='weather-model.tflite')
interpreter.allocate_tensors()

input_index = interpreter.get_input_details()[0]['index']
output_index = interpreter.get_output_details()[0]['index']

classes=['fogsmog',
         'hail',
         'lightning',
         'rain',
         'rainbow',
         'rime',
         'sandstorm',
         'snow']


#test URL
#url ='https://raw.githubusercontent.com/cjj1120/Weather_Image_Classification-Capstone_Project/main/Data/Add/Custom-test-img/rainbow-test.jpg'

def predict_url(url):
    img = download_image(url)
    img = prepare_image(img)
    x = np.array(img,dtype=np.float32)
    X = np.array([x])
    interpreter.set_tensor(input_index, X)
    interpreter.invoke()
    interpreter.get_tensor(output_index)
    pred = interpreter.get_tensor(output_index)
    float_pred = pred[0].tolist()
    dict_pred = dict(zip(classes, float_pred))
    return dict_pred

def lambda_handler(event, context):
    url = event['url']
    result = predict_url(url)
    return result