import torch
import cv2
from camere import get_picture
from contorl_motor import *
from predict import *
from PIL import Image
import RPi.GPIO  as GPIO

GPIO.setmode(GPIO.BCM)
GPIO.setwarnings(False)
channel = 18
GPIO.setup(channel, GPIO.IN, pull_up_down=GPIO.PUD_DOWN)
GPIO.add_event_detect(channel, GPIO.RISING, bouncetime=200)


def on_switch_pressed():
    print('开始分类')
    img = get_picture()
    image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # image = Image.open("test_image/img_15654.jpg")
    garbage_class = predict_garbage(image)
    contorl(garbage_class)


if __name__ == '__main__':
    print('请先放置垃圾')
    try:
        while True:
            if GPIO.event_detected(channel):
                on_switch_pressed()
                print('请先放置垃圾')
            time.sleep(0.01)

    except Exception as e:
        print(error)
    GPIO.cleanup()