Bugjudger
101 lines
4.0 KiB
Python
101 lines
4.0 KiB
Python
from os import path as osp
|
|
import os
|
|
import pickle
|
|
from PIL import Image
|
|
import imageio
|
|
from glob import glob
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
from sklearn.decomposition import PCA
|
|
from sklearn.preprocessing import minmax_scale
|
|
import cv2
|
|
|
|
|
|
def pca_image(data, mask, root, model_name, dataset, n_components=3, is_return=False):
|
|
features = data['embeddings']
|
|
ns,hw,c = features.shape
|
|
features = features.reshape(ns*hw,c)
|
|
mask = mask.reshape(ns*hw)
|
|
|
|
pca = PCA(n_components=n_components)
|
|
pca_features = pca.fit_transform(features[mask != 0])
|
|
pca_features = minmax_scale(pca_features, (0,255), axis=1)
|
|
# pca_features = minmax_scale(pca_features, (0,255), axis=0)
|
|
|
|
norm_features = np.zeros_like(mask,dtype=np.uint8).reshape(ns*hw,1).repeat(n_components,axis=1)
|
|
norm_features[mask != 0] = pca_features
|
|
|
|
if is_return:
|
|
norm_features = norm_features.reshape(1,ns,64,32,n_components)[...,:3].transpose(0,1,4,2,3) #
|
|
return norm_features
|
|
|
|
s = 20
|
|
assert ns % s == 0
|
|
norm_features = norm_features.reshape(ns//s,s,64,32,n_components)[...,:3].transpose(0,1,4,2,3)
|
|
data['embeddings'] = norm_features
|
|
save_image(data, root, model_name, dataset, need='image')
|
|
|
|
|
|
|
|
def save_image(data, root, model_name, dataset, need='image', mask=None):
|
|
images, label, seq_type, view = data['embeddings'], data['labels'], data['types'], data['views'] # n s c h w
|
|
if "image" in need:
|
|
root_path = os.path.join(root, dataset, model_name+'_image')
|
|
os.makedirs(os.path.join(root_path),exist_ok=True)
|
|
for i, id in enumerate(label[:]):
|
|
tmp = os.path.join(root_path, str(id).zfill(5), str(seq_type[i]), str(view[i]))
|
|
os.makedirs(tmp, exist_ok=True)
|
|
mb = None if mask is None else mask[i]
|
|
save_func(tmp, images[i], need, mb)
|
|
save_gif(tmp, tmp, str(view[i]))
|
|
|
|
if 'pkl' in need:
|
|
root_path = os.path.join(root, dataset, model_name+'_pkl')
|
|
os.makedirs(os.path.join(root_path),exist_ok=True)
|
|
for i, id in enumerate(label[:]):
|
|
tmp = os.path.join(root_path, str(id).zfill(5), str(seq_type[i]), str(view[i]))
|
|
os.makedirs(tmp, exist_ok=True)
|
|
mb = None if mask is None else mask[i]
|
|
save_func(tmp, images[i], 'pkl', mb)
|
|
|
|
if 'w' in need:
|
|
root_path = os.path.join(root, dataset, model_name+'_w')
|
|
os.makedirs(os.path.join(root_path),exist_ok=True)
|
|
for i, id in enumerate(label[:]):
|
|
tmp = os.path.join(root_path, str(id).zfill(5), str(seq_type[i]), str(view[i]))
|
|
os.makedirs(tmp, exist_ok=True)
|
|
mb = None if mask is None else mask[i]
|
|
save_func(tmp, data['w'], 'w', mb)
|
|
return
|
|
|
|
def save_func(tmp, data, ipts_type='image', mask=None):
|
|
if 'image' in ipts_type :
|
|
for i, con in enumerate(data):
|
|
if con.shape[0] == 1:
|
|
if 'jet' in ipts_type :
|
|
im = ((cv2.applyColorMap(con[0], cv2.COLORMAP_JET) * 0.5)[...,::-1] + 1.0*mask[i])
|
|
# im = mask[i]
|
|
im = np.clip(im,0,255).astype(np.uint8)
|
|
im = Image.fromarray(im, mode='RGB') # [h,w,c]
|
|
else:
|
|
im = Image.fromarray(con[0], mode='L')
|
|
else:
|
|
im = Image.fromarray(con.transpose(1,2,0), mode='RGB')
|
|
im.save(os.path.join(tmp, '%03d.png' % i))
|
|
elif ipts_type == 'pkl':
|
|
with open(os.path.join(tmp,'00.pkl'), 'wb') as f:
|
|
pickle.dump(data[:,0,:,:], f)
|
|
elif ipts_type == 'w':
|
|
for i in range(len(data)):
|
|
with open(os.path.join(tmp, str(i).zfill(2) + '.pkl'), 'wb') as f:
|
|
pickle.dump(data[i], f)
|
|
|
|
def save_gif(image_folder, save_folder, name="movie"):
|
|
images = []
|
|
filenames = sorted(glob(osp.join(image_folder, '*.png')))
|
|
# print(filenames)
|
|
for filename in filenames:
|
|
images.append(imageio.imread(filename))
|
|
imageio.mimsave(os.path.join(save_folder, f'{name}.gif'), images, duration=50, loop=0)
|