Update kaggle_sam2_gpu_1.py

kaggle_sam2_gpu_1.py

@@ -17,9 +17,9 @@ sys.path.append("/kaggle/florence-sam-kaggle")
 from utils.video import generate_unique_name, create_directory, delete_directory
 from utils.sam import load_sam_image_model, run_sam_inference, load_sam_video_model
 
-#DEVICE = torch.device("cuda")
+DEVICE = torch.device("cuda")
 DEVICE = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())][-1]
-#DEVICE = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())][0]
+DEVICE = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())][0]
 # DEVICE = torch.device("cpu")
 
 torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
@@ -32,7 +32,7 @@ SAM_VIDEO_MODEL = load_sam_video_model(device=DEVICE)
 VIDEO_SCALE_FACTOR = 1
 VIDEO_TARGET_DIRECTORY = "/kaggle/"
 create_directory(directory_path=VIDEO_TARGET_DIRECTORY)
-with open('/kaggle/output_video2.pkl', 'rb') as file:
+with open('/kaggle/output_video1.pkl', 'rb') as file:
     output_video = pickle.load(file)
 print(output_video)
 video_input= output_video
@@ -43,7 +43,7 @@ frame = next(frame_generator)
 frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
 
 
-with open('/kaggle/detections_list2.pkl', 'rb') as file:
+with open('/kaggle/detections_list1.pkl', 'rb') as file:
     detections_list = pickle.load(file)
 print(detections_list)
 detections = sv.Detections.merge(detections_list)
@@ -99,11 +99,11 @@ frames_generator = sv.get_video_frames_generator(video_input)
 masks_generator = SAM_VIDEO_MODEL.propagate_in_video(inference_state)
 n = 1
 
-out_dir = "/kaggle/output2"
+out_dir = "/kaggle/output1"
 os.makedirs(out_dir, exist_ok=True)
 import shutil
-shutil.rmtree('/kaggle/output2', ignore_errors=True)
-os.makedirs('/kaggle/output2', exist_ok=True)
+shutil.rmtree('/kaggle/output1', ignore_errors=True)
+os.makedirs('/kaggle/output1', exist_ok=True)
 with sv.VideoSink(video_path, video_info=video_info) as sink:
     for frame, (_, tracker_ids, mask_logits) in zip(frames_generator, masks_generator):
         frame = sv.scale_image(frame, VIDEO_SCALE_FACTOR)
@@ -146,7 +146,7 @@ with sv.VideoSink(video_path, video_info=video_info) as sink:
         frame = np.where(all_mask[:, :, None] > 0, masked_image, frame)
         # result 即为只保留 all_mask 遮罩内容的图像
         #frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        cv2.imwrite("/kaggle/output2/" + str(n) + ".jpeg", frame)
+        cv2.imwrite("/kaggle/output1/" + str(n) + ".jpeg", frame)
         n = n + 1