Update run.py

run.py

@@ -5,6 +5,14 @@ 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}
@@ -19,6 +27,7 @@ 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):
     """
@@ -31,14 +40,13 @@ def flip_periodically(im, interval_s=2):
     Returns:
         The flipped image.
     """
-    
-    # transforms.ToTensor()(im)
-    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
-    sleep(interval_s)
-    imout = (logits[0, 0] - logits[0, 0].min()) / (logits[0, 0].max() - logits[0, 0].min())
-    return imout  #np.flipud(im)
+    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)