Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import requests | |
| from io import BytesIO | |
| from transformers import AutoProcessor, CLIPSegForImageSegmentation | |
| from scipy.ndimage import label, find_objects | |
| import time | |
| # Streamlit 应用标题 | |
| st.title("使用图像分割模型分割证件照") | |
| # 输入图像 URL | |
| url = st.text_input("输入图像地址:", "https://i.ibb.co/GRCGQ3n/464.jpg") | |
| # 输入要识别的物体文本 | |
| texts_input = st.text_input("输入要检测的对象（以逗号分隔）:", "a card") | |
| texts = [text.strip() for text in texts_input.split(',')] | |
| # 选择面积阈值 | |
| area_threshold = st.slider("忽略小区域的面积阈值", 0, 10000, 5000) | |
| # 添加 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') | |
| st.write(f"设备: {device}") | |
| 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 | |
| # 标记连通区域 | |
| 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 | |
| # 仅保留原始图像中与当前矩形掩码匹配的区域 | |
| 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 # 计算运行时间 | |
| # 显示程序运行时间 | |
| st.write(f"程序运行时间: {elapsed_time:.2f} 秒") | |