Create 2.py

2.py

@@ -0,0 +1,172 @@
+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
+
+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'
+    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 = '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="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 = 'frame2'
+process_images_in_folder(input_folder)
+
+
+