import gradio as gr
import numpy as np
from time import sleep
import torch
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
# from torchvision import transforms

class Count:
    def __init__(self):
        self.n = 0
        self.imout = np.zeros((1000, 1000))
        
    def step(self):
        self.n += 1

        
weights2load = 'segformer_ep15_loss0.00.pth'
id2label = {0: 'seal', 255: 'bck'}
label2id = {'seal': 0, 'bck': 255}
model = SegformerForSemanticSegmentation.from_pretrained("nvidia/mit-b0",
                                                         num_labels=2,
                                                         id2label=id2label,
                                                         label2id=label2id,
)
image_processor = SegformerImageProcessor(reduce_labels=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.load_state_dict(torch.load(weights2load, weights_only=True, map_location=device))
model.to(device).eval()

counter = Count()

def flip_periodically(im, interval_s=2):
    """
    Flips the image periodically with the given interval.

    Args:
        im: The input image.
        interval_ms: The interval in milliseconds between flips.

    Returns:
        The flipped image.
    """
    counter.step()
    if (counter.imout.sum() == 0) or ((counter.n % 100) == 0):
        pixel_values = image_processor(im, return_tensors="pt").pixel_values.to(device)
        outputs = model(pixel_values=pixel_values)
        logits = outputs.logits.cpu().detach().numpy() ** 2
        counter.imout = (logits[0, 0] - logits[0, 0].min()) / (logits[0, 0].max() - logits[0, 0].min())
    return counter.imout, counter.n  #np.flipud(im)

with gr.Blocks() as demo:
    inp = gr.Image(sources=["webcam"], streaming=True)
    inp.stream(flip_periodically, inputs=inp, outputs=[gr.Image(), gr.Number()]) 


if __name__ == "__main__":

    demo.launch()