Spaces:
Paused
Paused
| import os | |
| import shutil | |
| import logging | |
| import tempfile | |
| import zipfile | |
| from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks | |
| from fastapi.responses import FileResponse, JSONResponse | |
| from typing import List | |
| import subprocess | |
| logging.basicConfig(level=logging.DEBUG, | |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) # '__name__' will use the module's name | |
| app = FastAPI() | |
| def remove_tmp(dir_path: str): | |
| # Cleanup: Remove the temporary directory and its contents | |
| shutil.rmtree(dir_path) | |
| async def upload_zip(background_tasks: BackgroundTasks, file: UploadFile = File(...)): | |
| if not file.filename.endswith('.zip'): | |
| raise HTTPException(status_code=400, detail="File must be a zip archive") | |
| tmpdir = tempfile.mkdtemp() | |
| try: | |
| # Save the zip file to the temporary directory | |
| zip_path = os.path.join(tmpdir, file.filename) | |
| with open(zip_path, 'wb') as buffer: | |
| shutil.copyfileobj(file.file, buffer) | |
| # Extract the zip file | |
| with zipfile.ZipFile(zip_path, 'r') as zip_ref: | |
| zip_ref.extractall(tmpdir) | |
| # Find all jpg files in the extracted directory | |
| image_paths = [os.path.join(tmpdir, f) for f in os.listdir(tmpdir) if f.endswith('.jpg')] | |
| if not image_paths: | |
| raise HTTPException(status_code=400, detail="No .jpg files found in the zip archive") | |
| # Prepare the output directory | |
| output_dir = os.path.join(tmpdir, "output") | |
| os.makedirs(output_dir, exist_ok=True) | |
| # Define the model weights path | |
| weights_path = "/app/dust3r/checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth" | |
| # # Check if CUDA is available and set the device | |
| # device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| # Run the 3D model generation script | |
| command = [ | |
| # "--silent", | |
| "python", "/app/dust3r/demo.py", | |
| "--images", *image_paths, | |
| "--weights", weights_path, | |
| # "--device", device, | |
| "--output_dir", output_dir | |
| ] | |
| try: | |
| result = subprocess.run(command, check=True, capture_output=True, text=True) | |
| logger.info("Subprocess output:", result.stdout) | |
| logger.warning("Subprocess error:", result.stderr) | |
| except subprocess.CalledProcessError as e: | |
| logger.error("Subprocess failed with error:", e.stderr) | |
| raise HTTPException(status_code=500, detail=f"3D model generation failed: {str(e)}") | |
| # Debug: List files in the output directory | |
| logger.debug("Files in output directory:", os.listdir(output_dir)) | |
| # Find the .glb file in the output directory | |
| glb_files = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if f.endswith('.glb')] | |
| if not glb_files: | |
| raise HTTPException(status_code=500, detail=f"No .glb file generated in {output_dir}") | |
| glb_file_path = glb_files[0] | |
| # Debug: Check if the file exists | |
| if not os.path.exists(glb_file_path): | |
| raise HTTPException(status_code=500, detail=f"File {glb_file_path} does not exist") | |
| # Optionally, remove the temporary directory if you don't need it anymore | |
| background_tasks.add_task(remove_tmp, tmpdir) | |
| # Return the .glb file as a response | |
| return FileResponse(glb_file_path, media_type='application/octet-stream', filename=os.path.basename(glb_file_path)) | |
| except Exception as e: | |
| logger.error("Error generating the scene.", exc_info=True) | |
| return JSONResponse({'error': 'Error generating the scene.', 'detail': str(e)}, status_code=500) | |
| if __name__ == '__main__': | |
| import uvicorn | |
| uvicorn.run(app, host='0.0.0.0', port=7860) | |