File size: 1,847 Bytes
1a3844d b41cba1 85f0e91 c596df4 86a4add 42159c7 f08406b 42159c7 c596df4 85f0e91 f08406b 85f0e91 b90cf97 b41cba1 1a3844d b41cba1 1a3844d b41cba1 f08406b 1a3844d b41cba1 1a3844d 85f0e91 1a3844d 85f0e91 b41cba1 |
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 59 |
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
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.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 #np.flipud(im)
with gr.Blocks() as demo:
inp = gr.Image(sources=["webcam"], streaming=True)
out = gr.Image()
inp.stream(flip_periodically, inputs=inp, outputs=out)
if __name__ == "__main__":
demo.launch() |