完善交互

app.py

@@ -3,10 +3,10 @@ import torch
 import numpy as np
 from PIL import Image
 import requests
-import matplotlib.pyplot as plt
 from io import BytesIO
 from transformers import AutoProcessor, CLIPSegForImageSegmentation
 from scipy.ndimage import label, find_objects
+import time
 
 # Streamlit 应用标题
 st.title("使用图像分割模型分割证件照")
@@ -21,72 +21,88 @@ texts = [text.strip() for text in texts_input.split(',')]
 # 选择面积阈值
 area_threshold = st.slider("忽略小区域的面积阈值", 0, 10000, 5000)
 
-# 加载模型和处理器
-processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
-model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
-
-# 下载并处理图像
-response = requests.get(url)
-image = Image.open(BytesIO(response.content))
-
-# 显示原始图像
-st.image(image, caption="原图", use_column_width=True)
-
-# 处理图像和文本
-inputs = processor(text=texts, images=[image] * len(texts), padding=True, return_tensors="pt")
-outputs = model(**inputs)
-
-# 将 logits 转换为概率值并生成掩码
-probabilities = torch.sigmoid(outputs.logits)
-masks = probabilities > 0.5
-masks = masks.detach().cpu().numpy()
-
-# 获取原始图像的 NumPy 数组
-image_np = np.array(image)
-
-# 全局计数器初始化
-global_counter = 1
-
-# 对每个物体生成分割图像
-for i, mask in enumerate(masks):
-    # 将掩码调整为与原始图像相同的尺寸
-    mask_resized = Image.fromarray(mask).resize((image_np.shape[1], image_np.shape[0]), resample=Image.LANCZOS)
-    mask_resized = np.array(mask_resized) > 0.5
-    
-    # 标记连通区域
-    labeled_mask, num_features = label(mask_resized)
-    object_slices = find_objects(labeled_mask)
-    
-    for j in range(1, num_features + 1):
-        # 获取当前连通区域的边界框
-        object_slice = object_slices[j-1]
-        area = (object_slice[0].stop - object_slice[0].start) * (object_slice[1].stop - object_slice[1].start)
-        
-        # 忽略面积小于阈值的区域
-        if area < area_threshold:
-            continue
-        
-        # 创建一个矩形掩码
-        single_object_mask = np.zeros_like(mask_resized)
-        single_object_mask[object_slice] = 1
+# 添加 GPU/CPU 选择按钮
+device_option = st.radio("选择设备", ("GPU", "CPU"))
+
+# 提交按钮
+if st.button('提交'):
+    # 在按钮点击后确定设备
+    device = torch.device('cuda' if device_option == 'GPU' and torch.cuda.is_available() else 'cpu')
+
+    start_time = time.time()  # 开始计时
+
+    # 加载模型和处理器到选定设备
+    processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
+    model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined").to(device)
+
+    # 下载并处理图像
+    response = requests.get(url)
+    image = Image.open(BytesIO(response.content))
+
+    # 显示原始图像
+    st.image(image, caption="原图", use_column_width=True)
+
+    # 处理图像和文本
+    inputs = processor(text=texts, images=[image] * len(texts), padding=True, return_tensors="pt").to(device)
+    outputs = model(**inputs)
+
+    # 将 logits 转换为概率值并生成掩码
+    probabilities = torch.sigmoid(outputs.logits)
+    masks = probabilities > 0.5
+    masks = masks.detach().cpu().numpy()  # 将数据移回 CPU 以进行后续处理
+
+    # 获取原始图像的 NumPy 数组
+    image_np = np.array(image)
+
+    # 全局计数器初始化
+    global_counter = 1
+
+    # 对每个物体生成分割图像
+    for i, mask in enumerate(masks):
+        # 将掩码调整为与原始图像相同的尺寸
+        mask_resized = Image.fromarray(mask).resize((image_np.shape[1], image_np.shape[0]), resample=Image.LANCZOS)
+        mask_resized = np.array(mask_resized) > 0.5
         
-        # 仅保留原始图像中与当前矩形掩码匹配的区域
-        single_segmented_image = np.zeros_like(image_np)
-        single_segmented_image[single_object_mask.astype(bool)] = image_np[single_object_mask.astype(bool)]
+        # 标记连通区域
+        labeled_mask, num_features = label(mask_resized)
+        object_slices = find_objects(labeled_mask)
         
-        # 去除黑色背景（即裁剪图像到非黑色区域）
-        non_black_area = np.any(single_segmented_image > 0, axis=-1)
-        if non_black_area.any():
-            rows = np.any(non_black_area, axis=1)
-            cols = np.any(non_black_area, axis=0)
-            rmin, rmax = np.where(rows)[0][[0, -1]]
-            cmin, cmax = np.where(cols)[0][[0, -1]]
+        for j in range(1, num_features + 1):
+            # 获取当前连通区域的边界框
+            object_slice = object_slices[j-1]
+            area = (object_slice[0].stop - object_slice[0].start) * (object_slice[1].stop - object_slice[1].start)
             
-            # 裁剪图像
-            cropped_image = single_segmented_image[rmin:rmax+1, cmin:cmax+1]
+            # 忽略面积小于阈值的区域
+            if area < area_threshold:
+                continue
             
-            # 显示裁剪后的图像
-            st.image(cropped_image, caption=f"{texts[i]} - 图 {global_counter}", use_column_width=True)
+            # 创建一个矩形掩码
+            single_object_mask = np.zeros_like(mask_resized)
+            single_object_mask[object_slice] = 1
+            
+            # 仅保留原始图像中与当前矩形掩码匹配的区域
+            single_segmented_image = np.zeros_like(image_np)
+            single_segmented_image[single_object_mask.astype(bool)] = image_np[single_object_mask.astype(bool)]
+            
+            # 去除黑色背景（即裁剪图像到非黑色区域）
+            non_black_area = np.any(single_segmented_image > 0, axis=-1)
+            if non_black_area.any():
+                rows = np.any(non_black_area, axis=1)
+                cols = np.any(non_black_area, axis=0)
+                rmin, rmax = np.where(rows)[0][[0, -1]]
+                cmin, cmax = np.where(cols)[0][[0, -1]]
+                
+                # 裁剪图像
+                cropped_image = single_segmented_image[rmin:rmax+1, cmin:cmax+1]
+                
+                # 显示裁剪后的图像
+                st.image(cropped_image, caption=f"{texts[i]} - 图 {global_counter}", use_column_width=True)
+
+                # 增加全局计数器
+                global_counter += 1
+
+    end_time = time.time()  # 结束计时
+    elapsed_time = end_time - start_time  # 计算运行时间
 
-            # 增加全局计数器
-            global_counter += 1
+    # 显示程序运行时间
+    st.write(f"程序运行时间: {elapsed_time:.2f} 秒")