jhj0517 commited on
Commit
8d01c02
1 Parent(s): c64775c

Update dtype

Browse files
Files changed (1) hide show
  1. modules/sam_inference.py +1 -2
modules/sam_inference.py CHANGED
@@ -40,7 +40,6 @@ class SamInference:
40
  self.current_model_type = DEFAULT_MODEL_TYPE
41
  self.model_dir = model_dir
42
  self.output_dir = output_dir
43
- self.model_path = os.path.join(self.model_dir, AVAILABLE_MODELS[DEFAULT_MODEL_TYPE][0])
44
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
45
  self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
46
  self.mask_generator = None
@@ -272,7 +271,7 @@ class SamInference:
272
  )
273
  images = get_frames_from_dir(vid_dir=TEMP_DIR, as_numpy=True)
274
 
275
- with torch.autocast(device_type=self.device, dtype=torch.float16):
276
  for out_frame_idx, out_obj_ids, out_mask_logits in generator:
277
  mask = (out_mask_logits[0] > 0.0).cpu().numpy()
278
  video_segments[out_frame_idx] = {
 
40
  self.current_model_type = DEFAULT_MODEL_TYPE
41
  self.model_dir = model_dir
42
  self.output_dir = output_dir
 
43
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
44
  self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
45
  self.mask_generator = None
 
271
  )
272
  images = get_frames_from_dir(vid_dir=TEMP_DIR, as_numpy=True)
273
 
274
+ with torch.autocast(device_type=self.device, dtype=self.dtype):
275
  for out_frame_idx, out_obj_ids, out_mask_logits in generator:
276
  mask = (out_mask_logits[0] > 0.0).cpu().numpy()
277
  video_segments[out_frame_idx] = {