78 lines
1.9 KiB
Python
78 lines
1.9 KiB
Python
import os.path
|
|
import shutil
|
|
|
|
|
|
from ytracking.track_ import *
|
|
from contrast.test_logic import group_image, inference
|
|
from tools.Interface import AiInterface, AiClass
|
|
|
|
from tools.initModel import models
|
|
import matplotlib.pyplot as plt
|
|
from sklearn.cluster import KMeans
|
|
|
|
models.initModel()
|
|
ai_obj = AiClass()
|
|
distance_lists = []
|
|
all_distance_lists = []
|
|
|
|
|
|
def showImg(newx):
|
|
plt.scatter(newx[:, 0], newx[:, 1])
|
|
plt.show()
|
|
|
|
|
|
def cosine_similarity(vec1, vec2):
|
|
vec1 = np.array(vec1)
|
|
vec2 = np.array(vec2)
|
|
cos_sim = vec1.dot(vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
|
|
return cos_sim
|
|
|
|
|
|
def get_kclusters(data, k):
|
|
Kmeans = KMeans(n_clusters=k, n_init='auto')
|
|
y_pred = Kmeans.fit_predict(data)
|
|
return y_pred
|
|
|
|
|
|
def move_file(pth, dirs, labels, y_pred):
|
|
for label, y_label in zip(labels, y_pred):
|
|
if not os.path.isdir(os.sep.join([pth, dirs, str(y_label)])):
|
|
os.mkdir(os.sep.join([pth, dirs, str(y_label)]))
|
|
try:
|
|
shutil.move(os.sep.join([pth, dirs, label]),
|
|
os.sep.join([pth, dirs, str(y_label), label]))
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
|
|
def createDic(pth, dirs):
|
|
all_dics = {}
|
|
imgs = os.sep.join([pth, dirs])
|
|
for img in os.listdir(imgs):
|
|
imgPth = os.sep.join([imgs, img])
|
|
feat_tensor = inference([imgPth], models.similarityModel, actionModel=False)
|
|
all_dics[img] = feat_tensor
|
|
return all_dics
|
|
|
|
|
|
def main(pth, dirs):
|
|
global distance_lists
|
|
allDics = createDic(pth, dirs)
|
|
labels = []
|
|
temp = None
|
|
for key, value in allDics.items():
|
|
value = value.cpu().detach().numpy()
|
|
labels.append(key)
|
|
if temp is None:
|
|
temp = value
|
|
else:
|
|
temp = np.vstack((temp, value))
|
|
y_pred = get_kclusters(temp, 3)
|
|
move_file(pth, dirs, labels, y_pred)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
pth = 'C:/Users/HP/Desktop/40084'
|
|
dirs = '1'
|
|
main(pth, dirs)
|