main.py

00001 import io
00002 import time
00003 from Classifier import ImageClassifier
00004 import numpy as np
00005 from PIL import Image
00006 from joblib import dump, load
00007 import picamera
00008 import serial
00009 
00010 
00011 
00012 print("Starting camera")
00013 camera = picamera.PiCamera()
00014 camera.resolution = (352, 240)
00015 camera.color_effects = (128, 128)  # turn camera to black and white
00016 
00017 mbed = serial.Serial('/dev/ttyACM0', baudrate=9600)
00018 time.sleep(1)
00019 
00020 print("starting model training")
00021 
00022 try:
00023     img_clf = load('model.joblib')
00024 except Exception as e:
00025     img_clf = ImageClassifier()
00026     # load images
00027     (train_raw, train_labels) = img_clf.load_data_from_folder('./train/')
00028     # convert images into features
00029     train_data = img_clf.extract_image_features(train_raw)
00030     # train model and test on training data
00031     img_clf.train_classifier(train_data, train_labels)
00032     # dump classifier into file for later use
00033     dump(img_clf, 'model.joblib')
00034 
00035 print("Model Trained")
00036 
00037 while True:
00038     stream = io.BytesIO()
00039     camera.capture(stream, format="bmp", resize=(352, 240))
00040     image = Image.open(stream)
00041     data = np.array(image)
00042     # Construct a numpy array from the stream
00043     features = img_clf.extract_image_features([data])
00044     label = img_clf.predict_labels(features)[0]
00045     if label == "disco":
00046         mbed.write("d")
00047     elif label == "tornado":
00048         mbed.write("t")
00049 
00050