import os from typing import Tuple, Optional import shutil import os import cv2 import numpy as np import spaces import supervision as sv import torch from PIL import Image from tqdm import tqdm from utils.video import generate_unique_name, create_directory, delete_directory from utils.florence import load_florence_model, run_florence_inference, \ FLORENCE_DETAILED_CAPTION_TASK, \ FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK, FLORENCE_OPEN_VOCABULARY_DETECTION_TASK from utils.modes import IMAGE_INFERENCE_MODES, IMAGE_OPEN_VOCABULARY_DETECTION_MODE, \ IMAGE_CAPTION_GROUNDING_MASKS_MODE, VIDEO_INFERENCE_MODES from utils.sam import load_sam_image_model, run_sam_inference, load_sam_video_model 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("cpu") torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__() if torch.cuda.get_device_properties(0).major >= 8: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE) SAM_IMAGE_MODEL = load_sam_image_model(device=DEVICE) SAM_VIDEO_MODEL = load_sam_video_model(device=DEVICE) # @title #视频帧提取 import supervision as sv import os import cv2 import shutil def extract_video_frames(video_input): # 目标目录 VIDEO_TARGET_DIRECTORY = '/kaggle/working/frame' if not os.path.exists(VIDEO_TARGET_DIRECTORY): os.makedirs(VIDEO_TARGET_DIRECTORY) shutil.rmtree(VIDEO_TARGET_DIRECTORY) # 视频缩放因子 VIDEO_SCALE_FACTOR = 1 # 获取视频信息 video_info = sv.VideoInfo.from_video_path(video_input) print(video_info) # 生成唯一的名称 # 使用视频文件名作为唯一名称 name = os.path.splitext(os.path.basename(video_input))[0] # 构建帧目录路径 frame_directory_path = os.path.join(VIDEO_TARGET_DIRECTORY, name) # 创建 ImageSink 对象 frames_sink = sv.ImageSink( target_dir_path=frame_directory_path, image_name_pattern="{:05d}.jpeg" ) # 获取视频帧生成器 frames_generator = sv.get_video_frames_generator(video_input) # 使用 with 语句确保资源正确释放 with frames_sink: # 遍历每一帧 for i, frame in enumerate(frames_generator): # 如果需要缩放帧 if VIDEO_SCALE_FACTOR != 1: frame = cv2.resize(frame, None, fx=VIDEO_SCALE_FACTOR, fy=VIDEO_SCALE_FACTOR) # 保存帧 frames_sink.save_image(frame) return frame_directory_path,video_info # 使用示例 video_input_path = '/kaggle/input/pinnpong/VS_010.mp4'# @param {type:"string"} video_frame_dir,video_info = extract_video_frames(video_input_path) texts = ['the table', 'all person','ball'] from PIL import Image import supervision as sv def detect_objects_in_image(image_input_path, texts): # 加载图像 image_input = Image.open(image_input_path) # 初始化检测列表 detections_list = [] # 对每个文本进行检测 for text in texts: _, result = run_florence_inference( model=FLORENCE_MODEL, processor=FLORENCE_PROCESSOR, device=DEVICE, image=image_input, task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK, text=text ) # 从结果中构建监督检测对象 detections = sv.Detections.from_lmm( lmm=sv.LMM.FLORENCE_2, result=result, resolution_wh=image_input.size ) # 运行 SAM 推理 detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections) # 将检测结果添加到列表中 detections_list.append(detections) # 合并所有检测结果 detections = sv.Detections.merge(detections_list) # 再次运行 SAM 推理 detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections) return detections # @title #合并遮罩加模糊merge_image_with_mask import numpy as np import cv2 import os from PIL import Image, ImageFilter def merge_image_with_mask(image_input_path, detections, output_folder): # 创建输出文件夹 if not os.path.exists(output_folder): os.makedirs(output_folder) # 提取图片文件名 image_name = os.path.basename(image_input_path) output_path = os.path.join(output_folder, image_name) # 创建掩码文件夹 mask_folder = '/kaggle/working/mask' if not os.path.exists(mask_folder): os.makedirs(mask_folder) # 合并掩码 combined_mask = np.zeros_like(detections.mask[0], dtype=np.uint8) for mask in detections.mask: combined_mask += mask combined_mask = np.clip(combined_mask, 0, 255) combined_mask = combined_mask.astype(np.uint8) # 膨胀掩码 kernel = np.ones((6, 6), np.uint8) dilated_mask = cv2.dilate(combined_mask, kernel, iterations=1) # 保存膨胀后的掩码 #mask_path = os.path.join(mask_folder, 'test1.png') #cv2.imwrite(mask_path, dilated_mask * 255) # 读取原始图像 original_image = cv2.imread(image_input_path) # 读取遮罩图片 #mask_image = cv2.imread(mask_path) # 确保原始图片和遮罩图片尺寸一致 #assert original_image.shape == mask_image.shape, "The images must have the same dimensions." # 使用掩膜从原始图片中提取部分区域 masked_image = cv2.bitwise_and(original_image, original_image, mask=dilated_mask) # 将掩膜应用于原始图片 blurred_image = cv2.GaussianBlur(original_image, (21, 21), 500) # 使用较大的核大小进行模糊 # 将提取的部分区域叠加到模糊后的图片上 blurred_image = cv2.bitwise_and(blurred_image, blurred_image, mask=~dilated_mask) # 将提取的部分区域叠加到模糊后的图片上 result = np.where(dilated_mask[:, :, None] > 0, masked_image, blurred_image) # 保存合并后的图片 cv2.imwrite(output_path, result) # @title #进度条批量处理文件夹process_images_in_folder(input_folder) from tqdm import tqdm import shutil def process_images_in_folder(input_folder): # 确保输出文件夹存在 output_folder = '/kaggle/working/okframe' if not os.path.exists(output_folder): os.makedirs(output_folder) shutil.rmtree('/kaggle/working/okframe') output_folder = '/kaggle/working/okframe' if not os.path.exists(output_folder): os.makedirs(output_folder) # 获取文件夹中的所有文件 files = [f for f in os.listdir(input_folder) if f.endswith('.jpg') or f.endswith('.png') or f.endswith('.jpeg')] # 使用 tqdm 显示进度条 for filename in tqdm(files, desc="Processing Images"): image_input_path = os.path.join(input_folder, filename) # 检测对象 detections = detect_objects_in_image( image_input_path=image_input_path, texts=texts ) # 合并图像 merge_image_with_mask( image_input_path=image_input_path, detections=detections, output_folder=output_folder ) # 使用示例 video_name = video_input_path.split('/')[-1].split('.')[0] input_folder = f'/kaggle/working/frame/{video_name}' process_images_in_folder(input_folder) # @title #合并所有帧成新视频frames_to_video(frame_folder, video_output_path, video_info) import cv2 import os import natsort import numpy as np def frames_to_video(frame_folder, video_output_path, video_info): # 获取所有帧文件名，并使用 natsorted 进行自然排序 frame_files = natsort.natsorted([f for f in os.listdir(frame_folder) if f.endswith(('.jpg', '.png', '.jpeg'))]) # 创建视频写入器 fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 编码器 out = cv2.VideoWriter(video_output_path, fourcc, video_info.fps, (video_info.width, video_info.height)) # 遍历所有帧文件 for frame_file in frame_files: frame_path = os.path.join(frame_folder, frame_file) frame = cv2.imread(frame_path) # 如果帧大小不匹配，调整大小 if frame.shape[:2] != (video_info.height, video_info.width): frame = cv2.resize(frame, (video_info.width, video_info.height)) # 写入视频 out.write(frame) # 释放资源 out.release() # 使用示例 video_info = video_info frame_folder = '/kaggle/working/okframe' video_output_path = '/kaggle/working/output_video.mp4' frames_to_video(frame_folder, video_output_path, video_info)