1.py

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)



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 = 'mask1'
if not os.path.exists(mask_folder):
    os.makedirs(mask_folder)
shutil.rmtree('mask1')
mask_folder = 'mask1'
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 = 'mask1'


    # 合并掩码
    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 = 'okframe1'
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    shutil.rmtree('okframe1')
    output_folder = 'okframe1'
    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 1 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 = 'frame1'
process_images_in_folder(input_folder)