Update run.py

run.py

@@ -17,7 +17,7 @@ from transformers import SegformerImageProcessor, SegformerForSemanticSegmentati
 #         self.n += 1
 
 
-# cnt = 0        
+cnt = 0  
 weights2load = 'segformer_ep15_loss0.00.pth'
 id2label = {0: 'seal', 255: 'bck'}
 label2id = {'seal': 0, 'bck': 255}
@@ -35,15 +35,19 @@ model.to(device).eval()
 # counter = Count()
 
 def segment(im, interval_s=2):
+    global cnt
+    cnt += 1
     im = cv2.resize(im, (im.shape[1] // 2, im.shape[0] // 2))
-    # if (counter.imout.sum() == 0) or ((cnt % 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
-    logits_n = (logits[0, 0] - logits[0, 0].min()) / (logits[0, 0].max() - logits[0, 0].min())
-    logits_n = cv2.resize(logits_n, (im.shape[1], im.shape[0]))
     imout = im.copy()
-    imout[..., 1] = np.clip(imout[..., 1] + logits_n * 200, 0, 254)
+    
+    if cnt % 10 == 0:
+        # if (counter.imout.sum() == 0) or ((cnt % 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
+        logits_n = (logits[0, 0] - logits[0, 0].min()) / (logits[0, 0].max() - logits[0, 0].min())
+        logits_n = cv2.resize(logits_n, (im.shape[1], im.shape[0]))
+        imout[..., 1] = np.clip(imout[..., 1] + logits_n * 200, 0, 254)
     return imout  #, cnt  #np.flipud(im)
 
 # with gr.Blocks() as demo: