2024-09-20 15:54:01 +08:00

178 lines
6.5 KiB
Python

from flask import Flask, request, jsonify, send_from_directory
import os
import json
import base64
import numpy as np
import pickle
from flask_cors import CORS
from data_load import DataLoadUtil
from reconstruction import ReconstructionUtil
from pts import PtsUtil
ROOT = os.path.join("./static")
print(ROOT)
app = Flask(__name__, static_folder="static")
CORS(app)
@app.route("/")
def serve_index():
return send_from_directory(app.static_folder, "index.html")
@app.route('/<path:filename>')
def serve_static(filename):
return send_from_directory(app.static_folder, filename)
@app.route('/get_scene_list', methods=['POST'])
def get_scene_list():
data = request.json
dataset_name = data.get('dataset_name')
dataset_path = os.path.join(ROOT, dataset_name)
if not os.path.exists(dataset_path):
print(f"Dataset not found: {dataset_path}")
return jsonify({"error": "Dataset not found"}), 404
scene_list = [d for d in os.listdir(dataset_path) if os.path.isdir(os.path.join(dataset_path, d))]
return jsonify({"scene_list": scene_list, "success": True})
@app.route('/get_scene_info', methods=['POST'])
def get_scene_info():
data = request.json
dataset_name = data.get('dataset_name')
scene_name = data.get('scene_name')
scene_path = os.path.join(ROOT, dataset_name, scene_name)
camera_params_path = os.path.join(scene_path, 'camera_params')
label_json_path = os.path.join(scene_path, 'label.json')
if not os.path.exists(scene_path) or not os.path.exists(label_json_path):
return jsonify({"error": "Scene or label.json not found"}), 404
with open(label_json_path, 'r') as f:
label_data = json.load(f)
sequence_length = len([f for f in os.listdir(camera_params_path) if os.path.isfile(os.path.join(camera_params_path, f))])
max_coverage_rate = label_data.get('max_coverage_rate')
best_sequence = label_data.get('best_sequence')
best_sequence_length = len(best_sequence)
best_sequence_formatted = []
for i in range(best_sequence_length):
best_sequence_formatted.append(
{
"frame": best_sequence[i][0],
"coverage_rate": round(best_sequence[i][1]*100,1)
}
)
return jsonify({
"sequence_length": sequence_length,
"max_coverage_rate": round(max_coverage_rate*100,2),
"best_sequence_length": best_sequence_length,
"best_sequence": best_sequence_formatted,
"success": True
})
def read_image_as_base64(file_path):
try:
with open(file_path, 'rb') as image_file:
encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
return encoded_string
except FileNotFoundError:
return None
@app.route('/get_frame_data', methods=['POST'])
def get_frame_data():
data = request.json
dataset_name = data.get('dataset_name')
scene_name = data.get('scene_name')
sequence = data.get('sequence')
scene_path = os.path.join(ROOT, dataset_name, scene_name)
root = os.path.join(ROOT, dataset_name)
camera_params_path = os.path.join(scene_path, 'camera_params')
depth_path = os.path.join(scene_path, 'depth')
mask_path = os.path.join(scene_path, 'mask')
model_points_normals = DataLoadUtil.load_points_normals(root, scene_name)
model_points = model_points_normals[:, :3]
obj_path = os.path.join(dataset_name, scene_name, 'mesh', 'world_target_mesh.obj')
mtl_path = os.path.join(dataset_name, scene_name, 'mesh', 'material.mtl')
if not all([os.path.exists(scene_path), os.path.exists(camera_params_path), os.path.exists(depth_path), os.path.exists(mask_path)]):
return jsonify({"error": "Invalid paths or files not found"}), 404
result = []
combined_point_cloud = np.zeros((0, 3))
last_CR = 0
for frame_info in sequence:
frame_id = frame_info.get('frame')
frame_data = {}
path = DataLoadUtil.get_path(root, scene_name, frame_id)
cam_params = DataLoadUtil.load_cam_info(path, binocular=True)
frame_data['cam_to_world'] = cam_params['cam_to_world'].tolist()
depth_file = os.path.join(depth_path, f'{frame_id}_L.png')
depth_base64 = read_image_as_base64(depth_file)
frame_data['depth'] = depth_base64 if depth_base64 else None
mask_file = os.path.join(mask_path, f'{frame_id}_L.png')
mask_base64 = read_image_as_base64(mask_file)
frame_data['mask'] = mask_base64 if mask_base64 else None
point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True)
sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_params['cam_to_world'], theta=75)
sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(sampled_point_cloud, 0.01)
frame_data['new_point_cloud'] = sampled_point_cloud.tolist()
frame_data['combined_point_cloud'] = combined_point_cloud.tolist()
combined_point_cloud = np.concatenate([combined_point_cloud, sampled_point_cloud], axis=0)
combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud, 0.01)
frame_data["coverage_rate"] = frame_info.get('coverage_rate')
delta_CR = frame_data["coverage_rate"] - last_CR
frame_data["delta_CR"] = round(delta_CR,2)
last_CR = frame_data["coverage_rate"]
result.append({
"frame_id": frame_id,
"data": frame_data
})
return jsonify({"seq_frame_data": result,"model_pts":model_points.tolist(), "obj_path": obj_path, "mtl_path":mtl_path, "success": True})
@app.route('/analysis_inference_result', methods=['POST'])
def analysis_inference_result():
res = {"success": True}
if 'file' not in request.files:
res["success"] = False
res["message"] = "No file part"
return jsonify(res)
file = request.files['file']
if file.filename == '':
res["success"] = False
res["message"] = "No selected file"
return jsonify(res)
try:
data = pickle.load(file)
except Exception as e:
res["success"] = False
res["message"] = f"File processing error: {e}"
return jsonify(res)
print(data)
return jsonify(res)
if __name__ == '__main__':
app.run(debug=True, port=13333)