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.florencegpu2 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("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) 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 mask_folder = 'mask2' if not os.path.exists(mask_folder): os.makedirs(mask_folder) shutil.rmtree('mask2') mask_folder = 'mask2' if not os.path.exists(mask_folder): os.makedirs(mask_folder) 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 = 'mask2' # 合并掩码 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, image_name) 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.medianBlur(original_image, 21) # 将提取的部分区域叠加到模糊后的图片上 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 = 'okframe2' if not os.path.exists(output_folder): os.makedirs(output_folder) shutil.rmtree('okframe2') output_folder = 'okframe2' 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="Gpu 2 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 ) # 使用示例 input_folder = 'frame2' process_images_in_folder(input_folder)