Update 1.py

1.py

@@ -1,3 +1,4 @@
+#%cd /workspace/florence-samflorence-sam
 import os
 from typing import Tuple, Optional
 import shutil
@@ -9,6 +10,11 @@ import supervision as sv
 import torch
 from PIL import Image
 from tqdm import tqdm
+import sys
+import json
+import pickle
+os.chdir("/workspace/florence-sam")
+sys.path.append('/workspace/florence-sam')
 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, \
@@ -19,161 +25,124 @@ from utils.sam import load_sam_image_model, run_sam_inference, load_sam_video_mo
 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
+with open('/workspace/florence-samtexts.pkl', 'rb') as file:
+    texts = pickle.load(file)
+print(texts)
+
+with open('/workspace/florence-samoutput_video.pkl', 'rb') as file:
+    output_video = pickle.load(file)
+print(output_video)
+
+VIDEO_SCALE_FACTOR = 1
+VIDEO_TARGET_DIRECTORY = "/workspace/florence-sam"
+create_directory(directory_path=VIDEO_TARGET_DIRECTORY)
+
+
+video_input= output_video
+texts = ['the table', 'men','ball']
+#VIDEO_TARGET_DIRECTORY = "/workspace/florence-sam"
+if not video_input:
+    print("Please upload a video.")
+'''
+frame_generator = sv.get_video_frames_generator(video_input)
+frame = next(frame_generator)
+frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+'''
+
+
+frame_generator = sv.get_video_frames_generator(video_input)
+
+# 获取视频的总帧数
+total_frames = int(sv.get_total_frames(video_input))
+
+# 计算中间帧的位置
+middle_frame_index = total_frames // 2
+with open('/workspace/florence-sam/middle_frame_index.pkl', 'wb') as file:
+    pickle.dump(middle_frame_index, file)
+# 读取到中间帧
+for _ in range(middle_frame_index):
+    frame = next(frame_generator)
+
+# 将帧从 BGR 转换为 RGB 并保存到 PIL 图像对象
+frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+
+detections_list = []
+width, height = frame.size
+all_ok_bboxes = []
+half_area = width * height * 0.5
+
+# 存储所有 the table 的边界框和面积
+table_bboxes = []
+table_areas = []
+given_area =1000
+ok_result =[]
+for text in texts:
+    _, result = run_florence_inference(
+        model=FLORENCE_MODEL,
+        processor=FLORENCE_PROCESSOR,
+        device=DEVICE,
+        image=frame,
+        task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK,
+        text=text    )
+    #print(result)
+    for bbox, label in zip(result['<OPEN_VOCABULARY_DETECTION>']['bboxes'], result['<OPEN_VOCABULARY_DETECTION>']['bboxes_labels']):
+      print(bbox, label)
+      new_result = {'<OPEN_VOCABULARY_DETECTION>': {'bboxes': [bbox], 'bboxes_labels': [label], 'polygons': [], 'polygons_labels': []}}
+      print(new_result)
+      if label == 'ping pong ball':
+          # 计算当前 ping pong ball 的面积
+          area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+          # 检查面积是否不超过给定边界框的面积
+          if area <= given_area:
+              all_ok_bboxes.append([[bbox[0], bbox[1]], [bbox[2], bbox[3]]])
+              ok_result.append(new_result)
+      elif label == 'the table':
+          # 计算当前 the table 的面积
+          print('the tablethe table!!!!')
+          area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
+          table_bboxes.append([[bbox[0] - 100, bbox[1]], [bbox[2] + 100, bbox[3]]])
+          table_areas.append(area)
+      elif label == 'table tennis bat':
+          all_ok_bboxes.append([[bbox[0], bbox[1]], [bbox[2], bbox[3]]])
+          ok_result.append(new_result)
+      elif label == 'men':
+          print('menmne!!!!')
+          all_ok_bboxes.append([[bbox[0], bbox[1]], [bbox[2], bbox[3]]])
+          ok_result.append(new_result)
+
+    # 找到面积最大的 the table
+    if table_areas:
+        max_area_index = table_areas.index(max(table_areas))
+        max_area_bbox = table_bboxes[max_area_index]
+        
+        # 检查面积是否超过50%
+        if max(table_areas) < half_area:
+            all_ok_bboxes.append(max_area_bbox)
+            ok_result.append(new_result)
+
+print(ok_result)
+with open('/workspace/florence-samall_ok_bboxes.pkl', 'wb') as file:
+    pickle.dump(all_ok_bboxes, file)
+
+for xyxy in ok_result:
+    print(frame.size,xyxy)
+    detections = sv.Detections.from_lmm(
+        lmm=sv.LMM.FLORENCE_2,
+        result=xyxy,
+        resolution_wh=frame.size
         )
-
-        # 从结果中构建监督检测对象
-        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)
-
-
-
+    detections = run_sam_inference(SAM_IMAGE_MODEL, frame, detections)
+    print(detections)
+    detections_list.append(detections)
+with open('/workspace/florence-samdetections_list.pkl', 'wb') as file:
+    pickle.dump(detections_list, file)
+print(detections_list)