Update run.py

run.py

@@ -7,16 +7,16 @@ from transformers import SegformerImageProcessor, SegformerForSemanticSegmentati
 
 
 
-class Count:
-    def __init__(self):
-        self.n = 0
-        self.imout = np.zeros((1000, 1000))
+# class Count:
+#     def __init__(self):
+#         self.n = 0
+#         self.imout = np.zeros((1000, 1000))
         
-    def step(self):
-        self.n += 1
+#     def step(self):
+#         self.n += 1
 
 
-cnt = 0        
+# cnt = 0        
 weights2load = 'segformer_ep15_loss0.00.pth'
 id2label = {0: 'seal', 255: 'bck'}
 label2id = {'seal': 0, 'bck': 255}
@@ -31,23 +31,25 @@ 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()
+# counter = Count()
 
-def flip_periodically(im, interval_s=2):
-    global cnt  # Access the global counter variable
-    cnt += 1
-    # counter.step()
-    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
-        counter.imout = (logits[0, 0] - logits[0, 0].min()) / (logits[0, 0].max() - logits[0, 0].min())
-    return counter.imout, cnt  #np.flipud(im)
+def segment(im, interval_s=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
+    imout = (logits[0, 0] - logits[0, 0].min()) / (logits[0, 0].max() - logits[0, 0].min())
+    return imout  #, cnt  #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()]) 
+# with gr.Blocks() as demo:
+#     inp = gr.Image(sources=["webcam"], streaming=True)
+#     inp.stream(segment, inputs=inp, outputs=[gr.Image(), gr.Number()]) 
 
+demo = gr.Interface(
+    segment,
+    [gr.Image(source="webcam", tool=None)],
+    ["image"],
+)
 
 if __name__ == "__main__":