93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
import torch
|
||
from config import config as conf
|
||
from PIL import Image
|
||
import numpy as np
|
||
|
||
|
||
def convert_rgba_to_rgb(image_path, output_path=None):
|
||
"""
|
||
将给定路径的4通道PNG图像转换为3通道，并保存到指定输出路径。
|
||
|
||
:param image_path: 输入图像的路径
|
||
:param output_path: 转换后的图像保存路径
|
||
"""
|
||
# 打开图像
|
||
img = Image.open(image_path)
|
||
# 转换图像模式从RGBA到RGB
|
||
# .convert('RGB')会丢弃Alpha通道并转换为纯RGB图像
|
||
if img.mode == 'RGBA':
|
||
# 转换为RGB模式
|
||
img_rgb = img.convert('RGB')
|
||
# 保存转换后的图像
|
||
img_rgb.save(image_path)
|
||
# print(f"Image converted from RGBA to RGB and saved to {image_path}")
|
||
# else:
|
||
# # 如果已经是RGB或其他模式，直接保存
|
||
# img.save(image_path)
|
||
# print(f"Image already in {img.mode} mode, saved to {image_path}")
|
||
|
||
|
||
def test_preprocess(images: list, actionModel=False) -> torch.Tensor:
|
||
res = []
|
||
for img in images:
|
||
try:
|
||
# print(img)
|
||
im = conf.test_transform(img) if actionModel else conf.test_transform(Image.open(img))
|
||
res.append(im)
|
||
except:
|
||
continue
|
||
data = torch.stack(res)
|
||
return data
|
||
|
||
|
||
def inference(images, model, actionModel=False):
|
||
data = test_preprocess(images, actionModel)
|
||
if torch.cuda.is_available():
|
||
data = data.to(conf.device)
|
||
features = model(data)
|
||
return features
|
||
|
||
|
||
def group_image(images, batch=64) -> list:
|
||
"""Group image paths by batch size"""
|
||
size = len(images)
|
||
res = []
|
||
for i in range(0, size, batch):
|
||
end = min(batch + i, size)
|
||
res.append(images[i:end])
|
||
return res
|
||
|
||
def normalize(queFeatList):
|
||
for num1 in range(len(queFeatList)):
|
||
for num2 in range(len(queFeatList[num1])):
|
||
queFeatList[num1][num2] = queFeatList[num1][num2] / np.linalg.norm(queFeatList[num1][num2])
|
||
return queFeatList
|
||
|
||
def getFeatureList(barList, imgList, model):
|
||
# featList = [[] for i in range(len(barList))]
|
||
# for index, feat in enumerate(imgList):
|
||
fe_nps = []
|
||
groups = group_image(imgList)
|
||
for group in groups:
|
||
feat_tensor = inference(group, model)
|
||
# for fe in feat_tensor:
|
||
if feat_tensor.device == 'cpu':
|
||
fe_np = feat_tensor.squeeze().detach().numpy()
|
||
# fe_np = fe_np[:, 256:]
|
||
# fe_np = fe_np.reshape(fe_np.shape[0], fe_np.shape[1], 1, 1)
|
||
else:
|
||
fe_np = feat_tensor.squeeze().detach().cpu().numpy()
|
||
# fe_np = fe_np[:, 256:]
|
||
# fe_np = fe_np[256:]
|
||
# fe_np = fe_np.reshape(fe_np.shape[0], fe_np.shape[1], 1, 1)
|
||
# fe_np = fe_np.reshape(1, fe_np.shape[0], 1, 1)
|
||
# print(fe_np)
|
||
|
||
fe_nps.append(fe_np)
|
||
# if fe_nps:
|
||
# merged_fe_np = np.concatenate(fe_nps, axis=0)
|
||
# else:
|
||
# merged_fe_np = np.array([]) #
|
||
# fe_list = normalize(fe_nps)
|
||
return fe_nps
|