-
Guillaume Duret authored5731489d
compute_features.py 9.63 KiB
import numpy as np
import json
from utils import compute_categories_id, compute_id_good_occ
from scipy.spatial.transform import Rotation
from bbox_2d import bbox_2d
import cv2
from instance_mask import instance
from pose import convert2
from matplotlib import image
from fps_alg import process2
import os
def process_compute(data_name, camera, camera_resized, new_size, Nb_camera, World_begin, Nb_world, list_categories, occ_target, vis):
transformation = np.matrix([[0.0000000, -1.0000000, 0.0000000],
[0.0000000, 0.0000000, -1.0000000],
[1.0000000, 0.0000000, 0.0000000]])
scenario = "Worlds"
destination_folders = [f"Generated_{scenario}_Testing", f"Generated_{scenario}_Evaluating", f"Generated_{scenario}_Training" ]
list_count_categories = {}
for destination_folder_loop in destination_folders : # [f"Generated_{scenario}_Testing", f"Generated_{scenario}_Evaluating", f"Generated_{scenario}_Training"] :
list_count_categories[destination_folder_loop] = {}
for i in range(World_begin, World_begin + Nb_world): # worlds
if i > 4 :
destination_folder = f"Generated_{scenario}_Testing"
elif i > 3 :
destination_folder = f"Generated_{scenario}_Evaluating"
else :
destination_folder = f"Generated_{scenario}_Training"
catergories_instance_array_id_to_cat, catergories_instance_array_cat_to_id, catergories_label_to_id = compute_categories_id(data_name, i)
for j in range(1, Nb_camera+1): # cameras
p = ((i-1)*Nb_camera) + j
catergories_occ_array = compute_id_good_occ(data_name, p, catergories_instance_array_id_to_cat, catergories_instance_array_cat_to_id, occ_target)
### 3D Poses ###
with open(f'{data_name}/Pose/{p}.json', 'r') as f:
data_3D_pose = json.load(f)
#print(data)
#print("len(data)", len(data_3D_pose))
### 2D BBox ###
with open(f"{data_name}/Bbox_2d/{p}.json", 'r') as f:
data_Bbox_2d = json.load(f)
with open(f"{data_name}/Bbox_3d/{p}.json", 'r') as f:
data_Bbox_3d = json.load(f)
if len(data_Bbox_2d) != len(data_3D_pose) :
raise TypeError("size of datas are differents !!")
if os.path.isfile(f'{data_name}/{destination_folder}/Count_{p-1}.json'):
with open(f'{data_name}/{destination_folder}/Count_{p-1}.json') as f:
list_count_categories[destination_folder] = json.load(f)
for categories in list_categories:
if categories in catergories_occ_array.keys():
Nb_instance = len(catergories_occ_array[categories])
meta = {}
if not categories in list_count_categories[destination_folder].keys():
#list_count_categories[categories] = {categories}
list_count_categories[destination_folder][categories] = {}
if Nb_instance in list_count_categories[destination_folder][categories].keys() :
list_count_categories[destination_folder][categories][Nb_instance] += 1
else :
list_count_categories[destination_folder][categories][Nb_instance] = 1
meta['id_generated'] = list_count_categories[destination_folder][categories][Nb_instance]
meta['id_original'] = p
meta['id_category'] = catergories_label_to_id[categories]
meta['id_instance'] = catergories_occ_array[categories]
meta['id_dataset'] = 1
meta['world'] = i
meta['camera'] = f"grabber_{j}"
meta['occlusion'] = occ_target
meta['Nb_instance_category'] = Nb_instance
if not os.path.isfile(f'{data_name}/{destination_folder}/{categories}/Meta_Gen/{categories}.json'):
with open(f'{data_name}/{destination_folder}/{categories}/Meta_Gen/{categories}.json', mode='w') as f:
feeds = {}
feeds[meta['id_generated']]=meta
f.write(json.dumps(feeds, indent=2))
else:
with open(f'{data_name}/{destination_folder}/{categories}/Meta_Gen/{categories}.json') as feedsjson:
feeds = json.load(feedsjson)
feeds[meta['id_generated']]=meta
with open(f'{data_name}/{destination_folder}/{categories}/Meta_Gen/{categories}.json', mode='w') as f:
f.write(json.dumps(feeds, indent=4))
if (Nb_instance == 1):
for k in range(len(data_3D_pose)):
if data_3D_pose[k]['id'] == catergories_occ_array[categories][0]:
rpy = data_3D_pose[k]['pose']['rpy']
rot = convert2(rpy)
R_exp = transformation @ rot
R_exp = np.array(R_exp)
xyz = data_3D_pose[k]['pose']['xyz']
T_exp = transformation @ xyz
T_exp = np.array(T_exp)
num_arr = np.c_[R_exp, T_exp[0]]
np.save(f'{data_name}/{destination_folder}/{categories}/Pose_transformed/{p}.npy', num_arr) # save
else:
continue
if data_Bbox_2d[k]['id'] == catergories_occ_array[categories][0]:
bbox = bbox_2d(data_Bbox_2d[k])
np.savetxt(f'{data_name}/{destination_folder}/{categories}/Bbox/{p}.txt', np.array(bbox).reshape((1, 4))) # save
else:
continue
if data_Bbox_3d[k]['id'] == catergories_occ_array[categories][0]:
bbox3d_size = data_Bbox_3d[k]['bbox']['size']
np.savetxt(f'{data_name}/{destination_folder}/{categories}/Bbox_3d_Gen/{p}.txt', bbox3d_size) # save
else:
continue
id = catergories_occ_array[categories][0]
img = cv2.imread(f"{data_name}/Instance_Segmentation/{p}.png", cv2.IMREAD_UNCHANGED) # plt.imread(path)
instance_img = instance(img, id)
cv2.imwrite(f"{data_name}/{destination_folder}/{categories}/Instance_Mask/{p}.png", 255*instance_img)
instance_img_resized = cv2.resize(instance_img, new_size)
cv2.imwrite(f"{data_name}/{destination_folder}/{categories}/Instance_Mask_resized/{p}.png", 255*instance_img_resized)
img = cv2.imread(f"{data_name}/RGB/{p}.png")
cv2.imwrite(f"{data_name}/{destination_folder}/{categories}/RGB_Gen/{p}.png", img)
img_resized = cv2.resize(img, new_size)
cv2.imwrite(f"{data_name}/{destination_folder}/{categories}/RGB_resized/{p}.png", img_resized)
np.set_printoptions(precision=15)
pose = np.load(f'{data_name}/{destination_folder}/{categories}/Pose_transformed/{p}.npy')
R_exp = pose[0:3, 0:3]
tVec = pose[0:3, 3]
fps_points = np.loadtxt(f'{data_name}/Generated/{categories}/{categories}_fps_3d.txt')
center = fps_points.mean(0)
fps_points = np.append(fps_points, [center], axis=0)
points = process2(fps_points, R_exp, tVec, camera, img, vis)
out = [int(catergories_occ_array[categories][0])] #len have to be 1 !!
ind = 1
for point in points:
x = point[0][0] / img.shape[1]
y = point[0][1] / img.shape[0]
out.append(x)
out.append(y)
ind += 2
np.savetxt(f'{data_name}/{destination_folder}/{categories}/FPS/{p}.txt', np.array(out).reshape(1, len(out)))
points_resized = process2(fps_points, R_exp, tVec, camera_resized, img_resized, vis)
out_resized = [int(catergories_occ_array[categories][0])] #len have to be 1 !
ind_resized = 1
for point_resized in points_resized:
x_resized = point_resized[0][0] / img_resized.shape[1]
y_resized = point_resized[0][1] / img_resized.shape[0]
out_resized.append(x_resized)
out_resized.append(y_resized)
ind_resized += 2
np.savetxt(f'{data_name}/{destination_folder}/{categories}/FPS_resized/{p}.txt', np.array(out_resized).reshape(1, len(out_resized)))
for destination_folder_loop in destination_folders : # [f"Generated_{scenario}_Testing", f"Generated_{scenario}_Evaluating", f"Generated_{scenario}_Training"] :
with open(f'{data_name}/{destination_folder_loop}/Count_{p}.json', mode='w') as f:
f.write(json.dumps(list_count_categories[destination_folder], indent=4))
print(list_count_categories[destination_folder])
print(list_count_categories)