from fastai.vision.all import *
import gradio as gr
import cv2

learn = load_learner('resnet18_model.pkl')

categories = ('Black', 'White')


def classify_image(image, webcam=None):
    if image is not None:
        is_people, _, probs = learn.predict(PILImage.create(image))
        return dict(zip(categories, map(float, probs)))
    elif webcam is not None:
        is_people, _, probs = learn.predict(PILImage.create(webcam))
        return dict(zip(categories, map(float, probs)))


def capture_image(cam):
    cam = cv2.VideoCapture(0)
    _, frame = cam.read()
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    cam.release()
    return frame


def webcam():
    cam = cv2.VideoCapture(0)
    while True:
        _, frame = cam.read()
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        yield frame


image = gr.inputs.Image(shape=(192, 192), label="Upload Image", source="upload")
webcam = gr.inputs.Image(shape=(192, 192), label="Webcam Image", source="webcam")
label = gr.outputs.Label()

examples = [['example1.jpg'], ['example2.jpg'], ['example3.jpg']]  # Nested list for multiple input components

intf = gr.Interface(fn=classify_image, inputs=[image, webcam], outputs=label, examples=examples,
                    allow_flagging='never')  # Set allow_flagging parameter as a string

intf.launch()