210 lines
6.9 KiB
Python
210 lines
6.9 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Created on Sun Nov 3 12:05:19 2024
|
||
|
||
@author: ym
|
||
"""
|
||
import os
|
||
import time
|
||
import torch
|
||
import pickle
|
||
import numpy as np
|
||
from config import config as conf
|
||
from model import resnet18 as resnet18
|
||
from feat_inference import inference_image
|
||
|
||
|
||
IMG_FORMAT = ['.bmp', '.jpg', '.jpeg', '.png']
|
||
|
||
'''======= 0. 配置特征提取模型地址 ======='''
|
||
model_path = conf.test_model
|
||
model_path = r"D:\exhibition\ckpt\zhanting.pth"
|
||
|
||
##============ load resnet mdoel
|
||
model = resnet18().to(conf.device)
|
||
# model = nn.DataParallel(model).to(conf.device)
|
||
model.load_state_dict(torch.load(model_path, map_location=conf.device))
|
||
model.eval()
|
||
print('load model {} '.format(conf.testbackbone))
|
||
|
||
def get_std_barcodeDict(bcdpath, savepath):
|
||
'''
|
||
inputs:
|
||
bcdpath: 已清洗的barcode样本图像，如果barcode下有'base'文件夹，只选用该文件夹下图像
|
||
(default = r'\\192.168.1.28\share\已标注数据备份\对比数据\barcode\barcode_1771')
|
||
功能：
|
||
生成并保存只有一个key值的字典 {barcode: [imgpath1, imgpath1, ...]}，
|
||
savepath: 字典存储地址，文件名格式：barcode.pickle
|
||
'''
|
||
|
||
# savepath = r'\\192.168.1.28\share\测试_202406\contrast\std_barcodes'
|
||
|
||
'''读取数据集中 barcode 列表'''
|
||
stdBarcodeList = []
|
||
for filename in os.listdir(bcdpath):
|
||
filepath = os.path.join(bcdpath, filename)
|
||
# if not os.path.isdir(filepath) or not filename.isdigit() or len(filename)<8:
|
||
# continue
|
||
stdBarcodeList.append(filename)
|
||
|
||
bcdPaths = [(barcode, os.path.join(bcdpath, barcode)) for barcode in stdBarcodeList]
|
||
|
||
'''遍历数据集，针对每一个barcode，生成并保存字典{barcode: [imgpath1, imgpath1, ...]}'''
|
||
k = 0
|
||
errbarcodes = []
|
||
for barcode, bpath in bcdPaths:
|
||
pickpath = os.path.join(savepath, f"{barcode}.pickle")
|
||
if os.path.isfile(pickpath):
|
||
continue
|
||
|
||
stdBarcodeDict = {}
|
||
stdBarcodeDict[barcode] = []
|
||
for root, dirs, files in os.walk(bpath):
|
||
imgpaths = []
|
||
if "base" in dirs:
|
||
broot = os.path.join(root, "base")
|
||
for imgname in os.listdir(broot):
|
||
imgpath = os.path.join(broot, imgname)
|
||
file, ext = os.path.splitext(imgpath)
|
||
|
||
if ext not in IMG_FORMAT:
|
||
continue
|
||
imgpaths.append(imgpath)
|
||
|
||
stdBarcodeDict[barcode].extend(imgpaths)
|
||
break
|
||
|
||
else:
|
||
for imgname in files:
|
||
imgpath = os.path.join(root, imgname)
|
||
_, ext = os.path.splitext(imgpath)
|
||
if ext not in IMG_FORMAT: continue
|
||
imgpaths.append(imgpath)
|
||
stdBarcodeDict[barcode].extend(imgpaths)
|
||
|
||
pickpath = os.path.join(savepath, f"{barcode}.pickle")
|
||
with open(pickpath, 'wb') as f:
|
||
pickle.dump(stdBarcodeDict, f)
|
||
print(f"Barcode: {barcode}")
|
||
|
||
# k += 1
|
||
# if k == 10:
|
||
# break
|
||
print(f"Len of errbarcodes: {len(errbarcodes)}")
|
||
return
|
||
|
||
|
||
|
||
def stdfeat_infer(imgPath, featPath, bcdSet=None):
|
||
'''
|
||
inputs:
|
||
imgPath: 该文件夹下的 pickle 文件格式 {barcode: [imgpath1, imgpath1, ...]}
|
||
featPath: imgPath图像对应特征的存储地址
|
||
功能:
|
||
对 imgPath中图像进行特征提取，生成只有一个key值的字典，
|
||
{barcode: features}，features.shape=(nsample, 256)，并保存至 featPath 中
|
||
|
||
'''
|
||
|
||
# imgPath = r"\\192.168.1.28\share\测试_202406\contrast\std_barcodes"
|
||
# featPath = r"\\192.168.1.28\share\测试_202406\contrast\std_features"
|
||
stdBarcodeDict = {}
|
||
stdBarcodeDict_ft16 = {}
|
||
|
||
|
||
'''4处同名: (1)barcode原始图像文件夹; (2)imgPath中的 .pickle 文件名、该pickle文件中字典的key值'''
|
||
|
||
k = 0
|
||
for filename in os.listdir(imgPath):
|
||
bcd, ext = os.path.splitext(filename)
|
||
pkpath = os.path.join(featPath, f"{bcd}.pickle")
|
||
|
||
if os.path.isfile(pkpath): continue
|
||
if bcdSet is not None and bcd not in bcdSet:
|
||
continue
|
||
|
||
filepath = os.path.join(imgPath, filename)
|
||
|
||
stdbDict = {}
|
||
stdbDict_ft16 = {}
|
||
stdbDict_uint8 = {}
|
||
|
||
t1 = time.time()
|
||
|
||
try:
|
||
with open(filepath, 'rb') as f:
|
||
bpDict = pickle.load(f)
|
||
for barcode, imgpaths in bpDict.items():
|
||
# feature = batch_inference(imgpaths, 8) #from vit distilled model of LiChen
|
||
feature = inference_image(imgpaths, conf.test_transform, model, conf.device)
|
||
feature /= np.linalg.norm(feature, axis=1)[:, None]
|
||
|
||
# float16
|
||
feature_ft16 = feature.astype(np.float16)
|
||
feature_ft16 /= np.linalg.norm(feature_ft16, axis=1)[:, None]
|
||
|
||
# uint8, 两种策略，1) 精度损失小, 2) 计算复杂度小
|
||
# feature_uint8, _ = ft16_to_uint8(feature_ft16)
|
||
feature_uint8 = (feature_ft16*128).astype(np.int8)
|
||
|
||
except Exception as e:
|
||
print(f"Error accured at: {filename}, with Exception is: {e}")
|
||
|
||
'''================ 保存单个barcode特征 ================'''
|
||
##================== float32
|
||
stdbDict["barcode"] = barcode
|
||
stdbDict["imgpaths"] = imgpaths
|
||
stdbDict["feats_ft32"] = feature
|
||
stdbDict["feats_ft16"] = feature_ft16
|
||
stdbDict["feats_uint8"] = feature_uint8
|
||
|
||
with open(pkpath, 'wb') as f:
|
||
pickle.dump(stdbDict, f)
|
||
|
||
stdBarcodeDict[barcode] = feature
|
||
stdBarcodeDict_ft16[barcode] = feature_ft16
|
||
|
||
t2 = time.time()
|
||
print(f"Barcode: {barcode}, need time: {t2-t1:.1f} secs")
|
||
# k += 1
|
||
# if k == 10:
|
||
# break
|
||
|
||
##================== float32
|
||
# pickpath = os.path.join(featPath, f"barcode_features_{k}.pickle")
|
||
# with open(pickpath, 'wb') as f:
|
||
# pickle.dump(stdBarcodeDict, f)
|
||
|
||
##================== float16
|
||
# pickpath_ft16 = os.path.join(featPath, f"barcode_features_ft16_{k}.pickle")
|
||
# with open(pickpath_ft16, 'wb') as f:
|
||
# pickle.dump(stdBarcodeDict_ft16, f)
|
||
|
||
return
|
||
|
||
|
||
|
||
def genfeatures(imgpath, bcdpath, featpath):
|
||
|
||
get_std_barcodeDict(imgpath, bcdpath)
|
||
stdfeat_infer(bcdpath, featpath, bcdSet=None)
|
||
|
||
print(f"Features have generated, saved in: {featpath}")
|
||
|
||
|
||
|
||
|
||
|
||
|
||
def main():
|
||
imgpath = r"\\192.168.1.28\share\展厅barcode数据\整理\zhantingBase"
|
||
bcdpath = r"D:\exhibition\dataset\bcdpath"
|
||
featpath = r"D:\exhibition\dataset\feats"
|
||
|
||
genfeatures(imgpath, bcdpath, featpath)
|
||
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|