image_to_3d_api.py

import os
import uuid
from flask import jsonify, send_file, request
from main import *
from PIL import Image
import torch
import numpy as np

def image_to_3d_func(image_path, output_path="output_3d.obj"):
    if image_to_3d_model is None:
        return "Image-to-3D model not initialized."
    pil_image = Image.open(image_path).convert("RGB")
    image = torch.tensor(np.array(pil_image)).float().permute(2,0,1).unsqueeze(0) / 255.0
    image = image.to(device)
    with torch.no_grad():
        mesh_obj = image_to_3d_model(image)
    with open(output_path, 'w') as f:
        f.write(mesh_obj)
    return output_path

def image_to_3d_api():
    if 'image' not in request.files:
        return jsonify({"error": "Image file is required"}), 400
    image_file = request.files['image']
    temp_image_path = f"temp_image_{uuid.uuid4()}.png"
    image_file.save(temp_image_path)
    output_file = image_to_3d_func(temp_image_path)
    os.remove(temp_image_path)
    if output_file == "Image-to-3D model not initialized.":
        return jsonify({"error": "Image to 3D failed"}), 500
    return send_file(output_file, mimetype="model/obj", as_attachment=True, download_name="output_3d.obj")