ai3d / app.py
TruthLens's picture
Update app.py
966572c verified
raw
history blame
3.07 kB
import os
import numpy as np
from PIL import Image
from stl import mesh
from flask import Flask, render_template_string, request, send_file
from werkzeug.utils import secure_filename
import cv2
app = Flask(__name__)
UPLOAD_FOLDER = 'uploads'
OUTPUT_FOLDER = 'outputs'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['OUTPUT_FOLDER'] = OUTPUT_FOLDER
HTML_TEMPLATE = """
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>3D Printing Model Generator</title>
</head>
<body>
<h1>Upload Image to Generate 3D Printable Model</h1>
<form action="/upload" method="post" enctype="multipart/form-data">
<label>Image:</label><input type="file" name="image" required><br><br>
<input type="submit" value="Upload & Generate 3D Model">
</form>
{% if download_link %}
<h2>3D Model Ready:</h2>
<a href="{{ download_link }}" download>Download STL File</a>
{% endif %}
</body>
</html>
"""
def image_to_height_map(image_path, size=(256, 256)):
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
resized = cv2.resize(image, size)
normalized = resized / 255.0
return normalized
def height_map_to_stl(height_map, output_path, scale=10.0):
rows, cols = height_map.shape
vertices = []
faces = []
for i in range(rows):
for j in range(cols):
z = height_map[i, j] * scale
vertices.append([j, i, z])
vertices = np.array(vertices).reshape((rows, cols, 3))
face_list = []
for i in range(rows - 1):
for j in range(cols - 1):
face_list.append([vertices[i, j], vertices[i + 1, j], vertices[i, j + 1]])
face_list.append([vertices[i + 1, j], vertices[i + 1, j + 1], vertices[i, j + 1]])
faces_np = np.array(face_list)
model = mesh.Mesh(np.zeros(faces_np.shape[0], dtype=mesh.Mesh.dtype))
for i, f in enumerate(faces_np):
model.vectors[i] = f
model.save(output_path)
def create_3d_model(image_path):
height_map = image_to_height_map(image_path)
output_path = os.path.join(app.config['OUTPUT_FOLDER'], 'model.stl')
height_map_to_stl(height_map, output_path)
return output_path
@app.route('/', methods=['GET'])
def index():
return render_template_string(HTML_TEMPLATE, download_link=None)
@app.route('/upload', methods=['POST'])
def upload():
uploaded_file = request.files['image']
if uploaded_file:
filename = secure_filename(uploaded_file.filename)
image_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
uploaded_file.save(image_path)
stl_path = create_3d_model(image_path)
return render_template_string(HTML_TEMPLATE, download_link=f'/download/{os.path.basename(stl_path)}')
return render_template_string(HTML_TEMPLATE, download_link=None)
@app.route('/download/<filename>', methods=['GET'])
def download_file(filename):
return send_file(os.path.join(app.config['OUTPUT_FOLDER'], filename), as_attachment=True)