noamholz commited on
Commit
457380a
·
verified ·
1 Parent(s): e9e0a79

Update run.py

Browse files
Files changed (1) hide show
  1. run.py +4 -2
run.py CHANGED
@@ -3,7 +3,7 @@ import numpy as np
3
  from time import sleep
4
  import torch
5
  from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
6
-
7
 
8
  weights2load = 'segformer_ep15_loss0.00.pth'
9
  id2label = {0: 'seal', 255: 'bck'}
@@ -31,7 +31,9 @@ def flip_periodically(im, interval_ms=0):
31
  Returns:
32
  The flipped image.
33
  """
34
- pixel_values = image_processor(im, return_tensors="pt").pixel_values.to(device)
 
 
35
  outputs = model(pixel_values=pixel_values)
36
  logits = outputs.logits.cpu()
37
  sleep(interval_ms / 1000) # Convert milliseconds to seconds
 
3
  from time import sleep
4
  import torch
5
  from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
6
+ from torchvision import transforms
7
 
8
  weights2load = 'segformer_ep15_loss0.00.pth'
9
  id2label = {0: 'seal', 255: 'bck'}
 
31
  Returns:
32
  The flipped image.
33
  """
34
+
35
+ transforms.ToTensor()(im)
36
+ image_processor(im, return_tensors="pt").pixel_values.to(device)
37
  outputs = model(pixel_values=pixel_values)
38
  logits = outputs.logits.cpu()
39
  sleep(interval_ms / 1000) # Convert milliseconds to seconds