mac9087 commited on
Commit
fa62b8d
·
verified ·
1 Parent(s): 352535f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +972 -125
app.py CHANGED
@@ -16,204 +16,1051 @@ from flask_cors import CORS
16
  import numpy as np
17
  import trimesh
18
  from transformers import pipeline
19
- from diffusers import StableDiffusionZero123Pipeline
20
- import imageio
21
- from scipy.spatial.transform import Rotation
22
 
23
  app = Flask(__name__)
24
- CORS(app)
25
 
26
- # Configuration
27
  UPLOAD_FOLDER = '/tmp/uploads'
28
  RESULTS_FOLDER = '/tmp/results'
29
  CACHE_DIR = '/tmp/huggingface'
30
  ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
31
- VIEW_ANGLES = [(30, 0), (30, 90), (30, 180), (30, 270)] # (elevation, azimuth)
32
 
 
33
  os.makedirs(UPLOAD_FOLDER, exist_ok=True)
34
  os.makedirs(RESULTS_FOLDER, exist_ok=True)
35
  os.makedirs(CACHE_DIR, exist_ok=True)
36
 
37
- # Environment variables for caching
38
  os.environ['HF_HOME'] = CACHE_DIR
39
  os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
 
40
 
41
  app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
42
- app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
43
 
44
- # Global models
45
- view_generator = None
 
 
46
  depth_estimator = None
47
  model_loaded = False
48
  model_loading = False
49
 
50
- processing_jobs = {}
 
 
51
 
 
52
  class TimeoutError(Exception):
53
  pass
54
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55
  def allowed_file(filename):
56
  return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
57
 
58
- def preprocess_image(image_path, size=256):
59
- img = Image.open(image_path).convert("RGB")
60
- img = img.resize((size, size), Image.LANCZOS)
61
- return img
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
62
 
63
- def load_models():
64
- global view_generator, depth_estimator, model_loaded
 
65
  if model_loaded:
66
- return
67
-
 
 
 
 
 
 
68
  try:
69
- # Load view generator
70
- view_generator = StableDiffusionZero123Pipeline.from_pretrained(
71
- "stabilityai/stable-zero123-6dof",
72
- torch_dtype=torch.float16,
73
- cache_dir=CACHE_DIR
74
- ).to("cuda" if torch.cuda.is_available() else "cpu")
75
-
76
- # Load depth estimator
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
77
  depth_estimator = pipeline(
78
- "depth-estimation",
79
- model="Intel/dpt-hybrid-midas",
 
80
  cache_dir=CACHE_DIR
81
  )
82
-
 
 
 
 
83
  model_loaded = True
84
- print("Models loaded successfully")
 
 
85
  except Exception as e:
86
- print(f"Error loading models: {str(e)}")
 
87
  raise
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
88
 
89
- def generate_novel_views(image, num_views=4):
90
- views = []
91
- for elevation, azimuth in VIEW_ANGLES:
92
- result = view_generator(
93
- image,
94
- num_inference_steps=50,
95
- elevation=elevation,
96
- azimuth=azimuth,
97
- guidance_scale=3.0
98
- ).images[0]
99
- views.append((result, (elevation, azimuth)))
100
- return views
101
-
102
- def depth_to_pointcloud(depth_map, pose, fov=60):
103
- h, w = depth_map.shape
104
- f = w / (2 * np.tan(np.radians(fov/2)))
105
-
106
- xx, yy = np.meshgrid(np.arange(w), np.arange(h))
107
- x = (xx - w/2) * depth_map / f
108
- y = (yy - h/2) * depth_map / f
109
- z = depth_map
110
-
111
- points = np.vstack((x.flatten(), y.flatten(), z.flatten())).T
112
-
113
- # Apply pose transformation
114
- rot = Rotation.from_euler('zyx', [pose[1], pose[0], 0], degrees=True)
115
- points = rot.apply(points)
116
-
117
- return points
118
-
119
- def create_mesh_from_pointcloud(points, image):
120
- pcd = trimesh.PointCloud(points)
121
- scene = pcd.scene()
122
- mesh = scene.delaunay_3d.triangulate_pcd(pcd)
123
- mesh.visual.vertex_colors = image.resize((mesh.vertices.shape[0], 3))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
124
  return mesh
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
125
 
126
  @app.route('/convert', methods=['POST'])
127
  def convert_image_to_3d():
 
128
  if 'image' not in request.files:
129
  return jsonify({"error": "No image provided"}), 400
130
 
131
  file = request.files['image']
 
 
 
132
  if not allowed_file(file.filename):
133
- return jsonify({"error": "Invalid file type"}), 400
134
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
135
  job_id = str(uuid.uuid4())
136
  output_dir = os.path.join(RESULTS_FOLDER, job_id)
137
  os.makedirs(output_dir, exist_ok=True)
138
-
 
139
  filename = secure_filename(file.filename)
140
  filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{job_id}_{filename}")
141
  file.save(filepath)
142
-
 
143
  processing_jobs[job_id] = {
144
  'status': 'processing',
145
  'progress': 0,
146
  'result_url': None,
147
- 'error': None
 
 
 
148
  }
149
-
 
150
  def process_image():
 
 
 
151
  try:
152
- # Preprocess input image
153
- img = preprocess_image(filepath)
154
- processing_jobs[job_id]['progress'] = 20
155
-
156
- # Generate novel views
157
- views = generate_novel_views(img)
158
- processing_jobs[job_id]['progress'] = 40
159
-
160
- # Process each view
161
- all_points = []
162
- for view_img, pose in views:
163
- # Estimate depth
164
- depth_result = depth_estimator(view_img)
165
- depth_map = np.array(depth_result["depth"])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
166
 
167
- # Convert to point cloud
168
- points = depth_to_pointcloud(depth_map, pose)
169
- all_points.append(points)
170
- processing_jobs[job_id]['progress'] += 10
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
171
 
172
- # Combine point clouds
173
- combined_points = np.vstack(all_points)
174
- processing_jobs[job_id]['progress'] = 80
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
175
 
176
- # Create mesh
177
- mesh = create_mesh_from_pointcloud(combined_points, img)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
178
 
179
- # Export
180
- obj_path = os.path.join(output_dir, "model.obj")
181
- mesh.export(obj_path)
182
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
183
  processing_jobs[job_id]['status'] = 'completed'
184
- processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
185
  processing_jobs[job_id]['progress'] = 100
186
-
 
 
 
 
 
 
 
 
 
 
 
187
  except Exception as e:
 
188
  processing_jobs[job_id]['status'] = 'error'
189
- processing_jobs[job_id]['error'] = str(e)
190
- finally:
 
191
  if os.path.exists(filepath):
192
  os.remove(filepath)
193
- gc.collect()
194
- torch.cuda.empty_cache()
195
-
196
- thread = threading.Thread(target=process_image)
197
- thread.start()
198
- return jsonify({"job_id": job_id}), 202
 
 
199
 
200
- @app.route('/download/<job_id>')
201
- def download_model(job_id):
202
  if job_id not in processing_jobs or processing_jobs[job_id]['status'] != 'completed':
203
- return jsonify({"error": "Job not complete"}), 404
204
 
205
- obj_path = os.path.join(RESULTS_FOLDER, job_id, "model.obj")
206
- return send_file(obj_path, as_attachment=True)
207
-
208
- @app.route('/progress/<job_id>')
209
- def get_progress(job_id):
210
- job = processing_jobs.get(job_id, {})
211
- return jsonify({
212
- 'status': job.get('status'),
213
- 'progress': job.get('progress'),
214
- 'error': job.get('error')
215
- })
 
 
 
216
 
217
  if __name__ == '__main__':
218
- load_models()
219
- app.run(host='0.0.0.0', port=7860)
 
 
 
 
 
16
  import numpy as np
17
  import trimesh
18
  from transformers import pipeline
19
+ from scipy.ndimage import gaussian_filter, uniform_filter, median_filter
20
+ from scipy import interpolate
21
+ import cv2
22
 
23
  app = Flask(__name__)
24
+ CORS(app) # Enable CORS for all routes
25
 
26
+ # Configure directories
27
  UPLOAD_FOLDER = '/tmp/uploads'
28
  RESULTS_FOLDER = '/tmp/results'
29
  CACHE_DIR = '/tmp/huggingface'
30
  ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
 
31
 
32
+ # Create necessary directories
33
  os.makedirs(UPLOAD_FOLDER, exist_ok=True)
34
  os.makedirs(RESULTS_FOLDER, exist_ok=True)
35
  os.makedirs(CACHE_DIR, exist_ok=True)
36
 
37
+ # Set Hugging Face cache environment variables
38
  os.environ['HF_HOME'] = CACHE_DIR
39
  os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
40
+ os.environ['HF_DATASETS_CACHE'] = os.path.join(CACHE_DIR, 'datasets')
41
 
42
  app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
43
+ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB max
44
 
45
+ # Job tracking dictionary
46
+ processing_jobs = {}
47
+
48
+ # Global model variables
49
  depth_estimator = None
50
  model_loaded = False
51
  model_loading = False
52
 
53
+ # Configuration for processing
54
+ TIMEOUT_SECONDS = 240 # 4 minutes max for processing
55
+ MAX_DIMENSION = 512 # Max image dimension to process
56
 
57
+ # TimeoutError for handling timeouts
58
  class TimeoutError(Exception):
59
  pass
60
 
61
+ # Thread-safe timeout implementation
62
+ def process_with_timeout(function, args, timeout):
63
+ result = [None]
64
+ error = [None]
65
+ completed = [False]
66
+
67
+ def target():
68
+ try:
69
+ result[0] = function(*args)
70
+ completed[0] = True
71
+ except Exception as e:
72
+ error[0] = e
73
+
74
+ thread = threading.Thread(target=target)
75
+ thread.daemon = True
76
+ thread.start()
77
+
78
+ thread.join(timeout)
79
+
80
+ if not completed[0]:
81
+ if thread.is_alive():
82
+ return None, TimeoutError(f"Processing timed out after {timeout} seconds")
83
+ elif error[0]:
84
+ return None, error[0]
85
+
86
+ if error[0]:
87
+ return None, error[0]
88
+
89
+ return result[0], None
90
+
91
  def allowed_file(filename):
92
  return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
93
 
94
+ # Enhanced image preprocessing with better detail preservation
95
+ def preprocess_image(image_path):
96
+ with Image.open(image_path) as img:
97
+ # Keep alpha channel if present
98
+ has_alpha = img.mode == 'RGBA'
99
+
100
+ # Convert to proper format while preserving alpha
101
+ if has_alpha:
102
+ img = img.convert("RGBA")
103
+ else:
104
+ img = img.convert("RGB")
105
+
106
+ # Resize if the image is too large
107
+ if img.width > MAX_DIMENSION or img.height > MAX_DIMENSION:
108
+ # Calculate new dimensions while preserving aspect ratio
109
+ if img.width > img.height:
110
+ new_width = MAX_DIMENSION
111
+ new_height = int(img.height * (MAX_DIMENSION / img.width))
112
+ else:
113
+ new_height = MAX_DIMENSION
114
+ new_width = int(img.width * (MAX_DIMENSION / img.height))
115
+
116
+ # Use high-quality Lanczos resampling for better detail preservation
117
+ img = img.resize((new_width, new_height), Image.LANCZOS)
118
+
119
+ # Convert to numpy array for additional preprocessing
120
+ img_array = np.array(img)
121
+
122
+ # Extract alpha channel if present
123
+ if has_alpha:
124
+ alpha = img_array[:, :, 3]
125
+ rgb = img_array[:, :, :3]
126
+ else:
127
+ rgb = img_array
128
+
129
+ # Apply adaptive histogram equalization for better contrast on RGB channels only
130
+ if len(rgb.shape) == 3 and rgb.shape[2] == 3:
131
+ # Convert to LAB color space for better contrast enhancement
132
+ lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB)
133
+ l, a, b = cv2.split(lab)
134
+
135
+ # Apply CLAHE to L channel
136
+ clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
137
+ cl = clahe.apply(l)
138
+
139
+ # Merge channels back
140
+ enhanced_lab = cv2.merge((cl, a, b))
141
+
142
+ # Convert back to RGB
143
+ rgb_enhanced = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2RGB)
144
+
145
+ # Recombine with alpha if needed
146
+ if has_alpha:
147
+ result = np.dstack((rgb_enhanced, alpha))
148
+ img = Image.fromarray(result, 'RGBA')
149
+ else:
150
+ img = Image.fromarray(rgb_enhanced, 'RGB')
151
+
152
+ return img
153
+
154
 
155
+ def load_model():
156
+ global depth_estimator, model_loaded, model_loading
157
+
158
  if model_loaded:
159
+ return depth_estimator
160
+
161
+ if model_loading:
162
+ # Wait for model to load if it's already in progress
163
+ while model_loading and not model_loaded:
164
+ time.sleep(0.5)
165
+ return depth_estimator
166
+
167
  try:
168
+ model_loading = True
169
+ print("Starting model loading...")
170
+
171
+ # Using DPT-Large which provides better detail than DPT-Hybrid
172
+ # Alternatively, consider "vinvino02/glpn-nyu" for different detail characteristics
173
+ model_name = "Intel/dpt-large"
174
+
175
+ # Download model with retry mechanism
176
+ max_retries = 3
177
+ retry_delay = 5
178
+
179
+ for attempt in range(max_retries):
180
+ try:
181
+ snapshot_download(
182
+ repo_id=model_name,
183
+ cache_dir=CACHE_DIR,
184
+ resume_download=True,
185
+ )
186
+ break
187
+ except Exception as e:
188
+ if attempt < max_retries - 1:
189
+ print(f"Download attempt {attempt+1} failed: {str(e)}. Retrying in {retry_delay} seconds...")
190
+ time.sleep(retry_delay)
191
+ retry_delay *= 2
192
+ else:
193
+ raise
194
+
195
+ # Initialize model with appropriate precision
196
+ device = "cuda" if torch.cuda.is_available() else "cpu"
197
+
198
+ # Load depth estimator pipeline
199
  depth_estimator = pipeline(
200
+ "depth-estimation",
201
+ model=model_name,
202
+ device=device if device == "cuda" else -1,
203
  cache_dir=CACHE_DIR
204
  )
205
+
206
+ # Optimize memory usage
207
+ if device == "cuda":
208
+ torch.cuda.empty_cache()
209
+
210
  model_loaded = True
211
+ print(f"Model loaded successfully on {device}")
212
+ return depth_estimator
213
+
214
  except Exception as e:
215
+ print(f"Error loading model: {str(e)}")
216
+ print(traceback.format_exc())
217
  raise
218
+ finally:
219
+ model_loading = False
220
+
221
+ # Enhanced depth processing function to improve detail quality
222
+ def enhance_depth_map(depth_map, detail_level='medium'):
223
+ """Apply sophisticated processing to enhance depth map details"""
224
+ # Convert to numpy array if needed
225
+ if isinstance(depth_map, Image.Image):
226
+ depth_map = np.array(depth_map)
227
+
228
+ # Make sure the depth map is 2D
229
+ if len(depth_map.shape) > 2:
230
+ depth_map = np.mean(depth_map, axis=2) if depth_map.shape[2] > 1 else depth_map[:,:,0]
231
+
232
+ # Create a copy for processing
233
+ enhanced_depth = depth_map.copy().astype(np.float32)
234
+
235
+ # Remove outliers using percentile clipping (more stable than min/max)
236
+ p_low, p_high = np.percentile(enhanced_depth, [1, 99])
237
+ enhanced_depth = np.clip(enhanced_depth, p_low, p_high)
238
+
239
+ # Normalize to 0-1 range for processing
240
+ enhanced_depth = (enhanced_depth - p_low) / (p_high - p_low) if p_high > p_low else enhanced_depth
241
+
242
+ # Apply different enhancement methods based on detail level
243
+ if detail_level == 'high':
244
+ # Apply unsharp masking for edge enhancement - simulating Hunyuan's detail technique
245
+ # First apply gaussian blur
246
+ blurred = gaussian_filter(enhanced_depth, sigma=1.5)
247
+ # Create the unsharp mask
248
+ mask = enhanced_depth - blurred
249
+ # Apply the mask with strength factor
250
+ enhanced_depth = enhanced_depth + 1.5 * mask
251
+
252
+ # Apply bilateral filter to preserve edges while smoothing noise
253
+ # Simulate using gaussian combinations
254
+ smooth1 = gaussian_filter(enhanced_depth, sigma=0.5)
255
+ smooth2 = gaussian_filter(enhanced_depth, sigma=2.0)
256
+ edge_mask = enhanced_depth - smooth2
257
+ enhanced_depth = smooth1 + 1.2 * edge_mask
258
+
259
+ elif detail_level == 'medium':
260
+ # Less aggressive but still effective enhancement
261
+ # Apply mild unsharp masking
262
+ blurred = gaussian_filter(enhanced_depth, sigma=1.0)
263
+ mask = enhanced_depth - blurred
264
+ enhanced_depth = enhanced_depth + 0.8 * mask
265
+
266
+ # Apply mild smoothing to reduce noise but preserve edges
267
+ enhanced_depth = gaussian_filter(enhanced_depth, sigma=0.5)
268
+
269
+ else: # low
270
+ # Just apply noise reduction without too much detail enhancement
271
+ enhanced_depth = gaussian_filter(enhanced_depth, sigma=0.7)
272
+
273
+ # Normalize again after processing
274
+ enhanced_depth = np.clip(enhanced_depth, 0, 1)
275
+
276
+ return enhanced_depth
277
 
278
+ # Convert depth map to 3D mesh with significantly enhanced detail
279
+ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
280
+ """Convert depth map to complete 3D model with all sides"""
281
+ # First, enhance the depth map for better details
282
+ enhanced_depth = enhance_depth_map(depth_map, detail_level)
283
+
284
+ # Get dimensions of depth map
285
+ h, w = enhanced_depth.shape
286
+
287
+ # Create a higher resolution grid for better detail
288
+ x = np.linspace(0, w-1, resolution)
289
+ y = np.linspace(0, h-1, resolution)
290
+ x_grid, y_grid = np.meshgrid(x, y)
291
+
292
+ # Use bicubic interpolation for smoother surface
293
+ interp_func = interpolate.RectBivariateSpline(
294
+ np.arange(h), np.arange(w), enhanced_depth, kx=3, ky=3
295
+ )
296
+
297
+ # Sample depth at grid points
298
+ z_values = interp_func(y, x, grid=True)
299
+
300
+ # Process enhancement as in original code
301
+ if detail_level == 'high':
302
+ dx = np.gradient(z_values, axis=1)
303
+ dy = np.gradient(z_values, axis=0)
304
+ gradient_magnitude = np.sqrt(dx**2 + dy**2)
305
+ edge_mask = np.clip(gradient_magnitude * 5, 0, 0.2)
306
+ z_values = z_values + edge_mask * (z_values - gaussian_filter(z_values, sigma=1.0))
307
+
308
+ # Normalize z-values with advanced scaling
309
+ z_min, z_max = np.percentile(z_values, [2, 98])
310
+ z_values = (z_values - z_min) / (z_max - z_min) if z_max > z_min else z_values
311
+
312
+ # Apply depth scaling
313
+ if detail_level == 'high':
314
+ z_scaling = 2.5
315
+ elif detail_level == 'medium':
316
+ z_scaling = 2.0
317
+ else:
318
+ z_scaling = 1.5
319
+
320
+ z_values = z_values * z_scaling
321
+
322
+ # Normalize coordinates for front face
323
+ x_grid_front = (x_grid / w - 0.5) * 2.0
324
+ y_grid_front = (y_grid / h - 0.5) * 2.0
325
+
326
+ # Create all vertices (front, back, and sides)
327
+ vertices = []
328
+
329
+ # Front face vertices
330
+ front_vertices = np.vstack([x_grid_front.flatten(), -y_grid_front.flatten(), -z_values.flatten()]).T
331
+ vertices.append(front_vertices)
332
+
333
+ # Back face vertices (mirrored from front face)
334
+ back_depth = 1.0 # Constant thickness for the model
335
+ back_vertices = np.vstack([x_grid_front.flatten(), -y_grid_front.flatten(), -z_values.flatten() - back_depth]).T
336
+ vertices.append(back_vertices)
337
+
338
+ # Create side vertices (top, bottom, left, right)
339
+ # For simplicity, we use a grid mapping for sides
340
+ top_vertices = []
341
+ bottom_vertices = []
342
+ left_vertices = []
343
+ right_vertices = []
344
+
345
+ # Create sides by connecting front and back faces
346
+ for i in range(resolution):
347
+ # Top edge
348
+ for j in range(resolution):
349
+ if i == 0:
350
+ top_vertices.append(front_vertices[i * resolution + j])
351
+ top_vertices.append(back_vertices[i * resolution + j])
352
+ # Bottom edge
353
+ if i == resolution - 1:
354
+ bottom_vertices.append(front_vertices[i * resolution + j])
355
+ bottom_vertices.append(back_vertices[i * resolution + j])
356
+ # Left edge
357
+ if j == 0:
358
+ left_vertices.append(front_vertices[i * resolution + j])
359
+ left_vertices.append(back_vertices[i * resolution + j])
360
+ # Right edge
361
+ if j == resolution - 1:
362
+ right_vertices.append(front_vertices[i * resolution + j])
363
+ right_vertices.append(back_vertices[i * resolution + j])
364
+
365
+ # Combine all vertices
366
+ all_vertices = np.vstack([
367
+ front_vertices,
368
+ back_vertices,
369
+ np.array(top_vertices),
370
+ np.array(bottom_vertices),
371
+ np.array(left_vertices),
372
+ np.array(right_vertices)
373
+ ])
374
+
375
+ # Create faces (triangles)
376
+ faces = []
377
+
378
+ # Front face triangles
379
+ front_faces = []
380
+ for i in range(resolution-1):
381
+ for j in range(resolution-1):
382
+ p1 = i * resolution + j
383
+ p2 = i * resolution + (j + 1)
384
+ p3 = (i + 1) * resolution + j
385
+ p4 = (i + 1) * resolution + (j + 1)
386
+
387
+ # Calculate normals for consistent orientation
388
+ v1 = front_vertices[p1]
389
+ v2 = front_vertices[p2]
390
+ v3 = front_vertices[p3]
391
+ v4 = front_vertices[p4]
392
+
393
+ norm1 = np.cross(v2-v1, v4-v1)
394
+ norm2 = np.cross(v4-v3, v1-v3)
395
+
396
+ if np.dot(norm1, norm2) >= 0:
397
+ front_faces.append([p1, p2, p4])
398
+ front_faces.append([p1, p4, p3])
399
+ else:
400
+ front_faces.append([p1, p2, p3])
401
+ front_faces.append([p2, p4, p3])
402
+
403
+ # Back face triangles (note: reversed winding order for correct normals)
404
+ back_offset = resolution * resolution # Offset for back face vertices
405
+ back_faces = []
406
+ for i in range(resolution-1):
407
+ for j in range(resolution-1):
408
+ p1 = back_offset + i * resolution + j
409
+ p2 = back_offset + i * resolution + (j + 1)
410
+ p3 = back_offset + (i + 1) * resolution + j
411
+ p4 = back_offset + (i + 1) * resolution + (j + 1)
412
+
413
+ # Reverse winding order compared to front face
414
+ back_faces.append([p1, p4, p2])
415
+ back_faces.append([p1, p3, p4])
416
+
417
+ # Side faces (connecting front and back)
418
+ side_faces = []
419
+
420
+ # Add faces for sides (top, bottom, left, right)
421
+ side_offset = 2 * resolution * resolution # Offset after front and back
422
+
423
+ # Top side
424
+ top_count = len(top_vertices)
425
+ for i in range(0, top_count - 2, 2):
426
+ side_faces.append([side_offset + i, side_offset + i + 1, side_offset + i + 3])
427
+ side_faces.append([side_offset + i, side_offset + i + 3, side_offset + i + 2])
428
+
429
+ # Bottom side
430
+ bottom_offset = side_offset + top_count
431
+ bottom_count = len(bottom_vertices)
432
+ for i in range(0, bottom_count - 2, 2):
433
+ side_faces.append([bottom_offset + i, bottom_offset + i + 3, bottom_offset + i + 1])
434
+ side_faces.append([bottom_offset + i, bottom_offset + i + 2, bottom_offset + i + 3])
435
+
436
+ # Left side
437
+ left_offset = bottom_offset + bottom_count
438
+ left_count = len(left_vertices)
439
+ for i in range(0, left_count - 2, 2):
440
+ side_faces.append([left_offset + i, left_offset + i + 1, left_offset + i + 3])
441
+ side_faces.append([left_offset + i, left_offset + i + 3, left_offset + i + 2])
442
+
443
+ # Right side
444
+ right_offset = left_offset + left_count
445
+ right_count = len(right_vertices)
446
+ for i in range(0, right_count - 2, 2):
447
+ side_faces.append([right_offset + i, right_offset + i + 3, right_offset + i + 1])
448
+ side_faces.append([right_offset + i, right_offset + i + 2, right_offset + i + 3])
449
+
450
+ # Combine all faces
451
+ faces = np.array(front_faces + back_faces + side_faces)
452
+
453
+ # Create mesh
454
+ mesh = trimesh.Trimesh(vertices=all_vertices, faces=faces)
455
+
456
+ # Apply texturing if image is provided
457
+ if image:
458
+ # Handle RGBA properly to ensure transparency is maintained
459
+ img_array = np.array(image)
460
+
461
+ # Check if image has alpha channel
462
+ has_alpha = len(img_array.shape) == 3 and img_array.shape[2] == 4
463
+
464
+ # Create vertex colors with transparency support
465
+ vertex_colors = np.zeros((all_vertices.shape[0], 4), dtype=np.uint8)
466
+
467
+ # Fill with default color (will be overridden for front face)
468
+ vertex_colors[:, :3] = [200, 200, 200] # Light gray default
469
+ vertex_colors[:, 3] = 255 # Fully opaque
470
+
471
+ # Front face texture (sample from image)
472
+ for i in range(resolution):
473
+ for j in range(resolution):
474
+ # Calculate image coordinates
475
+ img_x = j * (img_array.shape[1] - 1) / (resolution - 1)
476
+ img_y = i * (img_array.shape[0] - 1) / (resolution - 1)
477
+
478
+ # Bilinear interpolation setup
479
+ x0, y0 = int(img_x), int(img_y)
480
+ x1, y1 = min(x0 + 1, img_array.shape[1] - 1), min(y0 + 1, img_array.shape[0] - 1)
481
+
482
+ # Interpolation weights
483
+ wx = img_x - x0
484
+ wy = img_y - y0
485
+
486
+ vertex_idx = i * resolution + j
487
+
488
+ if has_alpha:
489
+ # Handle RGBA with bilinear interpolation
490
+ for c in range(4):
491
+ vertex_colors[vertex_idx, c] = int((1-wx)*(1-wy)*img_array[y0, x0, c] +
492
+ wx*(1-wy)*img_array[y0, x1, c] +
493
+ (1-wx)*wy*img_array[y1, x0, c] +
494
+ wx*wy*img_array[y1, x1, c])
495
+ else:
496
+ # Handle RGB (no alpha)
497
+ for c in range(3):
498
+ vertex_colors[vertex_idx, c] = int((1-wx)*(1-wy)*img_array[y0, x0, c] +
499
+ wx*(1-wy)*img_array[y0, x1, c] +
500
+ (1-wx)*wy*img_array[y1, x0, c] +
501
+ wx*wy*img_array[y1, x1, c])
502
+ vertex_colors[vertex_idx, 3] = 255 # Fully opaque
503
+
504
+ # Apply simpler texturing to back face
505
+ back_face_start = resolution * resolution
506
+ back_face_color = [180, 180, 180, 255] # Slightly darker gray
507
+ vertex_colors[back_face_start:back_face_start + (resolution * resolution)] = back_face_color
508
+
509
+ mesh.visual.vertex_colors = vertex_colors
510
+
511
+ # Apply smoothing to get rid of staircase artifacts
512
+ if detail_level != 'high':
513
+ mesh = mesh.smoothed(method='laplacian', iterations=1)
514
+
515
+ # Calculate and fix normals for better rendering
516
+ mesh.fix_normals()
517
+
518
  return mesh
519
+
520
+
521
+
522
+ @app.route('/health', methods=['GET'])
523
+ def health_check():
524
+ return jsonify({
525
+ "status": "healthy",
526
+ "model": "Enhanced Depth-Based 3D Model Generator (DPT-Large)",
527
+ "device": "cuda" if torch.cuda.is_available() else "cpu"
528
+ }), 200
529
+
530
+ @app.route('/progress/<job_id>', methods=['GET'])
531
+ def progress(job_id):
532
+ def generate():
533
+ if job_id not in processing_jobs:
534
+ yield f"data: {json.dumps({'error': 'Job not found'})}\n\n"
535
+ return
536
+
537
+ job = processing_jobs[job_id]
538
+
539
+ # Send initial progress
540
+ yield f"data: {json.dumps({'status': 'processing', 'progress': job['progress']})}\n\n"
541
+
542
+ # Wait for job to complete or update
543
+ last_progress = job['progress']
544
+ check_count = 0
545
+ while job['status'] == 'processing':
546
+ if job['progress'] != last_progress:
547
+ yield f"data: {json.dumps({'status': 'processing', 'progress': job['progress']})}\n\n"
548
+ last_progress = job['progress']
549
+
550
+ time.sleep(0.5)
551
+ check_count += 1
552
+
553
+ # If client hasn't received updates for a while, check if job is still running
554
+ if check_count > 60: # 30 seconds with no updates
555
+ if 'thread_alive' in job and not job['thread_alive']():
556
+ job['status'] = 'error'
557
+ job['error'] = 'Processing thread died unexpectedly'
558
+ break
559
+ check_count = 0
560
+
561
+ # Send final status
562
+ if job['status'] == 'completed':
563
+ yield f"data: {json.dumps({'status': 'completed', 'progress': 100, 'result_url': job['result_url'], 'preview_url': job['preview_url']})}\n\n"
564
+ else:
565
+ yield f"data: {json.dumps({'status': 'error', 'error': job['error']})}\n\n"
566
+
567
+ return Response(stream_with_context(generate()), mimetype='text/event-stream')
568
 
569
  @app.route('/convert', methods=['POST'])
570
  def convert_image_to_3d():
571
+ # Check if image is in the request
572
  if 'image' not in request.files:
573
  return jsonify({"error": "No image provided"}), 400
574
 
575
  file = request.files['image']
576
+ if file.filename == '':
577
+ return jsonify({"error": "No image selected"}), 400
578
+
579
  if not allowed_file(file.filename):
580
+ return jsonify({"error": f"File type not allowed. Supported types: {', '.join(ALLOWED_EXTENSIONS)}"}), 400
581
+
582
+ # Get optional parameters with defaults
583
+ try:
584
+ mesh_resolution = min(int(request.form.get('mesh_resolution', 100)), 200) # Limit max resolution
585
+ output_format = request.form.get('output_format', 'obj').lower()
586
+ detail_level = request.form.get('detail_level', 'medium').lower() # Parameter for detail level
587
+ texture_quality = request.form.get('texture_quality', 'medium').lower() # New parameter for texture quality
588
+ except ValueError:
589
+ return jsonify({"error": "Invalid parameter values"}), 400
590
+
591
+ # Validate output format
592
+ if output_format not in ['obj', 'glb']:
593
+ return jsonify({"error": "Unsupported output format. Use 'obj' or 'glb'"}), 400
594
+
595
+ # Adjust mesh resolution based on detail level
596
+ if detail_level == 'high':
597
+ mesh_resolution = min(int(mesh_resolution * 1.5), 200)
598
+ elif detail_level == 'low':
599
+ mesh_resolution = max(int(mesh_resolution * 0.7), 50)
600
+
601
+ # Create a job ID
602
  job_id = str(uuid.uuid4())
603
  output_dir = os.path.join(RESULTS_FOLDER, job_id)
604
  os.makedirs(output_dir, exist_ok=True)
605
+
606
+ # Save the uploaded file
607
  filename = secure_filename(file.filename)
608
  filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{job_id}_{filename}")
609
  file.save(filepath)
610
+
611
+ # Initialize job tracking
612
  processing_jobs[job_id] = {
613
  'status': 'processing',
614
  'progress': 0,
615
  'result_url': None,
616
+ 'preview_url': None,
617
+ 'error': None,
618
+ 'output_format': output_format,
619
+ 'created_at': time.time()
620
  }
621
+
622
+ # Start processing in a separate thread
623
  def process_image():
624
+ thread = threading.current_thread()
625
+ processing_jobs[job_id]['thread_alive'] = lambda: thread.is_alive()
626
+
627
  try:
628
+ # Preprocess image with enhanced detail preservation
629
+ processing_jobs[job_id]['progress'] = 5
630
+ image = preprocess_image(filepath)
631
+ processing_jobs[job_id]['progress'] = 10
632
+
633
+ # Load model
634
+ try:
635
+ model = load_model()
636
+ processing_jobs[job_id]['progress'] = 30
637
+ except Exception as e:
638
+ processing_jobs[job_id]['status'] = 'error'
639
+ processing_jobs[job_id]['error'] = f"Error loading model: {str(e)}"
640
+ return
641
+
642
+ # Process image with thread-safe timeout
643
+ try:
644
+ def estimate_depth():
645
+ # Get depth map
646
+ result = model(image)
647
+ depth_map = result["depth"]
648
+
649
+ # Convert to numpy array if needed
650
+ if isinstance(depth_map, torch.Tensor):
651
+ depth_map = depth_map.cpu().numpy()
652
+ elif hasattr(depth_map, 'numpy'):
653
+ depth_map = depth_map.numpy()
654
+ elif isinstance(depth_map, Image.Image):
655
+ depth_map = np.array(depth_map)
656
+
657
+ return depth_map
658
 
659
+ depth_map, error = process_with_timeout(estimate_depth, [], TIMEOUT_SECONDS)
660
+
661
+ if error:
662
+ if isinstance(error, TimeoutError):
663
+ processing_jobs[job_id]['status'] = 'error'
664
+ processing_jobs[job_id]['error'] = f"Processing timed out after {TIMEOUT_SECONDS} seconds"
665
+ return
666
+ else:
667
+ raise error
668
+
669
+ processing_jobs[job_id]['progress'] = 60
670
+
671
+ # Create mesh from depth map with enhanced detail handling
672
+ mesh_resolution_int = int(mesh_resolution)
673
+ mesh = depth_to_mesh(depth_map, image, resolution=mesh_resolution_int, detail_level=detail_level)
674
+ processing_jobs[job_id]['progress'] = 80
675
+
676
+ except Exception as e:
677
+ error_details = traceback.format_exc()
678
+ processing_jobs[job_id]['status'] = 'error'
679
+ processing_jobs[job_id]['error'] = f"Error during processing: {str(e)}"
680
+ print(f"Error processing job {job_id}: {str(e)}")
681
+ print(error_details)
682
+ return
683
+
684
+ # Export based on requested format with enhanced quality settings
685
+ try:
686
+ if output_format == 'obj':
687
+ obj_path = os.path.join(output_dir, "model.obj")
688
+
689
+ # Export with normal and texture coordinates
690
+ mesh.export(
691
+ obj_path,
692
+ file_type='obj',
693
+ include_normals=True,
694
+ include_texture=True
695
+ )
696
+
697
+ # Create a zip file with OBJ and MTL
698
+ zip_path = os.path.join(output_dir, "model.zip")
699
+ with zipfile.ZipFile(zip_path, 'w') as zipf:
700
+ zipf.write(obj_path, arcname="model.obj")
701
+ mtl_path = os.path.join(output_dir, "model.mtl")
702
+ if os.path.exists(mtl_path):
703
+ zipf.write(mtl_path, arcname="model.mtl")
704
+
705
+ # Include texture file if it exists
706
+ texture_path = os.path.join(output_dir, "model.png")
707
+ if os.path.exists(texture_path):
708
+ zipf.write(texture_path, arcname="model.png")
709
+
710
+ processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
711
+ processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
712
+
713
+ elif output_format == 'glb':
714
+ # Export as GLB with enhanced settings
715
+ glb_path = os.path.join(output_dir, "model.glb")
716
+ mesh.export(
717
+ glb_path,
718
+ file_type='glb'
719
+ )
720
+
721
+ processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
722
+ processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
723
+
724
+ # Update job status
725
+ processing_jobs[job_id]['status'] = 'completed'
726
+ processing_jobs[job_id]['progress'] = 100
727
+ print(f"Job {job_id} completed successfully")
728
+ except Exception as e:
729
+ error_details = traceback.format_exc()
730
+ processing_jobs[job_id]['status'] = 'error'
731
+ processing_jobs[job_id]['error'] = f"Error exporting model: {str(e)}"
732
+ print(f"Error exporting model for job {job_id}: {str(e)}")
733
+ print(error_details)
734
+
735
+ # Clean up temporary file
736
+ if os.path.exists(filepath):
737
+ os.remove(filepath)
738
+
739
+ # Force garbage collection to free memory
740
+ gc.collect()
741
+ if torch.cuda.is_available():
742
+ torch.cuda.empty_cache()
743
+
744
+ except Exception as e:
745
+ # Handle errors
746
+ error_details = traceback.format_exc()
747
+ processing_jobs[job_id]['status'] = 'error'
748
+ processing_jobs[job_id]['error'] = f"{str(e)}\n{error_details}"
749
+ print(f"Error processing job {job_id}: {str(e)}")
750
+ print(error_details)
751
+
752
+ # Clean up on error
753
+ if os.path.exists(filepath):
754
+ os.remove(filepath)
755
+
756
+ # Start processing thread
757
+ processing_thread = threading.Thread(target=process_image)
758
+ processing_thread.daemon = True
759
+ processing_thread.start()
760
+
761
+ # Return job ID immediately
762
+ return jsonify({"job_id": job_id}), 202 # 202 Accepted
763
 
764
+ @app.route('/download/<job_id>', methods=['GET'])
765
+ def download_model(job_id):
766
+ if job_id not in processing_jobs or processing_jobs[job_id]['status'] != 'completed':
767
+ return jsonify({"error": "Model not found or processing not complete"}), 404
768
+
769
+ # Get the output directory for this job
770
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
771
+
772
+ # Determine file format from the job data
773
+ output_format = processing_jobs[job_id].get('output_format', 'obj')
774
+
775
+ if output_format == 'obj':
776
+ zip_path = os.path.join(output_dir, "model.zip")
777
+ if os.path.exists(zip_path):
778
+ return send_file(zip_path, as_attachment=True, download_name="model.zip")
779
+ else: # glb
780
+ glb_path = os.path.join(output_dir, "model.glb")
781
+ if os.path.exists(glb_path):
782
+ return send_file(glb_path, as_attachment=True, download_name="model.glb")
783
+
784
+ return jsonify({"error": "File not found"}), 404
785
 
786
+ @app.route('/preview/<job_id>', methods=['GET'])
787
+ def preview_model(job_id):
788
+ if job_id not in processing_jobs or processing_jobs[job_id]['status'] != 'completed':
789
+ return jsonify({"error": "Model not found or processing not complete"}), 404
790
+
791
+ # Get the output directory for this job
792
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
793
+ output_format = processing_jobs[job_id].get('output_format', 'obj')
794
+
795
+ if output_format == 'obj':
796
+ obj_path = os.path.join(output_dir, "model.obj")
797
+ if os.path.exists(obj_path):
798
+ return send_file(obj_path, mimetype='model/obj')
799
+ else: # glb
800
+ glb_path = os.path.join(output_dir, "model.glb")
801
+ if os.path.exists(glb_path):
802
+ return send_file(glb_path, mimetype='model/gltf-binary')
803
+
804
+ return jsonify({"error": "Model file not found"}), 404
805
+
806
+ # Cleanup old jobs periodically
807
+ def cleanup_old_jobs():
808
+ current_time = time.time()
809
+ job_ids_to_remove = []
810
+
811
+ for job_id, job_data in processing_jobs.items():
812
+ # Remove completed jobs after 1 hour
813
+ if job_data['status'] == 'completed' and (current_time - job_data.get('created_at', 0)) > 3600:
814
+ job_ids_to_remove.append(job_id)
815
+ # Remove error jobs after 30 minutes
816
+ elif job_data['status'] == 'error' and (current_time - job_data.get('created_at', 0)) > 1800:
817
+ job_ids_to_remove.append(job_id)
818
+
819
+ # Remove the jobs
820
+ for job_id in job_ids_to_remove:
821
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
822
+ try:
823
+ import shutil
824
+ if os.path.exists(output_dir):
825
+ shutil.rmtree(output_dir)
826
+ except Exception as e:
827
+ print(f"Error cleaning up job {job_id}: {str(e)}")
828
+
829
+ # Remove from tracking dictionary
830
+ if job_id in processing_jobs:
831
+ del processing_jobs[job_id]
832
+
833
+ # Schedule the next cleanup
834
+ threading.Timer(300, cleanup_old_jobs).start() # Run every 5 minutes
835
+
836
+ # New endpoint to get detailed information about a model
837
+ @app.route('/model-info/<job_id>', methods=['GET'])
838
+ def model_info(job_id):
839
+ if job_id not in processing_jobs:
840
+ return jsonify({"error": "Model not found"}), 404
841
+
842
+ job = processing_jobs[job_id]
843
+
844
+ if job['status'] != 'completed':
845
+ return jsonify({
846
+ "status": job['status'],
847
+ "progress": job['progress'],
848
+ "error": job.get('error')
849
+ }), 200
850
+
851
+ # For completed jobs, include information about the model
852
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
853
+ model_stats = {}
854
+
855
+ # Get file size
856
+ if job['output_format'] == 'obj':
857
+ obj_path = os.path.join(output_dir, "model.obj")
858
+ zip_path = os.path.join(output_dir, "model.zip")
859
+
860
+ if os.path.exists(obj_path):
861
+ model_stats['obj_size'] = os.path.getsize(obj_path)
862
 
863
+ if os.path.exists(zip_path):
864
+ model_stats['package_size'] = os.path.getsize(zip_path)
 
865
 
866
+ else: # glb
867
+ glb_path = os.path.join(output_dir, "model.glb")
868
+ if os.path.exists(glb_path):
869
+ model_stats['model_size'] = os.path.getsize(glb_path)
870
+
871
+ # Return detailed info
872
+ return jsonify({
873
+ "status": job['status'],
874
+ "model_format": job['output_format'],
875
+ "download_url": job['result_url'],
876
+ "preview_url": job['preview_url'],
877
+ "model_stats": model_stats,
878
+ "created_at": job.get('created_at'),
879
+ "completed_at": job.get('completed_at')
880
+ }), 200
881
+
882
+ @app.route('/', methods=['GET'])
883
+ def index():
884
+ return jsonify({
885
+ "message": "Enhanced Image to 3D API (DPT-Large Model)",
886
+ "endpoints": [
887
+ "/convert",
888
+ "/progress/<job_id>",
889
+ "/download/<job_id>",
890
+ "/preview/<job_id>",
891
+ "/model-info/<job_id>"
892
+ ],
893
+ "parameters": {
894
+ "mesh_resolution": "Integer (50-200), controls mesh density",
895
+ "output_format": "obj or glb",
896
+ "detail_level": "low, medium, or high - controls the level of detail in the final model",
897
+ "texture_quality": "low, medium, or high - controls the quality of textures"
898
+ },
899
+ "description": "This API creates high-quality 3D models from 2D images with enhanced detail finishing similar to Hunyuan model"
900
+ }), 200
901
+
902
+ # Example endpoint showing how to compare different detail levels
903
+ @app.route('/detail-comparison', methods=['POST'])
904
+ def compare_detail_levels():
905
+ # Check if image is in the request
906
+ if 'image' not in request.files:
907
+ return jsonify({"error": "No image provided"}), 400
908
+
909
+ file = request.files['image']
910
+ if file.filename == '':
911
+ return jsonify({"error": "No image selected"}), 400
912
+
913
+ if not allowed_file(file.filename):
914
+ return jsonify({"error": f"File type not allowed. Supported types: {', '.join(ALLOWED_EXTENSIONS)}"}), 400
915
+
916
+ # Create a job ID
917
+ job_id = str(uuid.uuid4())
918
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
919
+ os.makedirs(output_dir, exist_ok=True)
920
+
921
+ # Save the uploaded file
922
+ filename = secure_filename(file.filename)
923
+ filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{job_id}_{filename}")
924
+ file.save(filepath)
925
+
926
+ # Initialize job tracking
927
+ processing_jobs[job_id] = {
928
+ 'status': 'processing',
929
+ 'progress': 0,
930
+ 'result_url': None,
931
+ 'preview_url': None,
932
+ 'error': None,
933
+ 'output_format': 'glb', # Use GLB for comparison
934
+ 'created_at': time.time(),
935
+ 'comparison': True
936
+ }
937
+
938
+ # Process in separate thread to create 3 different detail levels
939
+ def process_comparison():
940
+ thread = threading.current_thread()
941
+ processing_jobs[job_id]['thread_alive'] = lambda: thread.is_alive()
942
+
943
+ try:
944
+ # Preprocess image
945
+ image = preprocess_image(filepath)
946
+ processing_jobs[job_id]['progress'] = 10
947
+
948
+ # Load model
949
+ try:
950
+ model = load_model()
951
+ processing_jobs[job_id]['progress'] = 20
952
+ except Exception as e:
953
+ processing_jobs[job_id]['status'] = 'error'
954
+ processing_jobs[job_id]['error'] = f"Error loading model: {str(e)}"
955
+ return
956
+
957
+ # Process image to get depth map
958
+ try:
959
+ depth_map = model(image)["depth"]
960
+ if isinstance(depth_map, torch.Tensor):
961
+ depth_map = depth_map.cpu().numpy()
962
+ elif hasattr(depth_map, 'numpy'):
963
+ depth_map = depth_map.numpy()
964
+ elif isinstance(depth_map, Image.Image):
965
+ depth_map = np.array(depth_map)
966
+
967
+ processing_jobs[job_id]['progress'] = 40
968
+ except Exception as e:
969
+ processing_jobs[job_id]['status'] = 'error'
970
+ processing_jobs[job_id]['error'] = f"Error estimating depth: {str(e)}"
971
+ return
972
+
973
+ # Create meshes at different detail levels
974
+ result_urls = {}
975
+
976
+ for detail_level in ['low', 'medium', 'high']:
977
+ try:
978
+ # Update progress
979
+ if detail_level == 'low':
980
+ processing_jobs[job_id]['progress'] = 50
981
+ elif detail_level == 'medium':
982
+ processing_jobs[job_id]['progress'] = 70
983
+ else:
984
+ processing_jobs[job_id]['progress'] = 90
985
+
986
+ # Create mesh with appropriate detail level
987
+ mesh_resolution = 100 # Fixed resolution for fair comparison
988
+ if detail_level == 'high':
989
+ mesh_resolution = 150
990
+ elif detail_level == 'low':
991
+ mesh_resolution = 80
992
+
993
+ mesh = depth_to_mesh(depth_map, image,
994
+ resolution=mesh_resolution,
995
+ detail_level=detail_level)
996
+
997
+ # Export as GLB
998
+ model_path = os.path.join(output_dir, f"model_{detail_level}.glb")
999
+ mesh.export(model_path, file_type='glb')
1000
+
1001
+ # Add to result URLs
1002
+ result_urls[detail_level] = f"/compare-download/{job_id}/{detail_level}"
1003
+
1004
+ except Exception as e:
1005
+ print(f"Error processing {detail_level} detail level: {str(e)}")
1006
+ # Continue with other detail levels even if one fails
1007
+
1008
+ # Update job status
1009
  processing_jobs[job_id]['status'] = 'completed'
 
1010
  processing_jobs[job_id]['progress'] = 100
1011
+ processing_jobs[job_id]['result_urls'] = result_urls
1012
+ processing_jobs[job_id]['completed_at'] = time.time()
1013
+
1014
+ # Clean up temporary file
1015
+ if os.path.exists(filepath):
1016
+ os.remove(filepath)
1017
+
1018
+ # Force garbage collection
1019
+ gc.collect()
1020
+ if torch.cuda.is_available():
1021
+ torch.cuda.empty_cache()
1022
+
1023
  except Exception as e:
1024
+ # Handle errors
1025
  processing_jobs[job_id]['status'] = 'error'
1026
+ processing_jobs[job_id]['error'] = f"Error during processing: {str(e)}"
1027
+
1028
+ # Clean up on error
1029
  if os.path.exists(filepath):
1030
  os.remove(filepath)
1031
+
1032
+ # Start processing thread
1033
+ processing_thread = threading.Thread(target=process_comparison)
1034
+ processing_thread.daemon = True
1035
+ processing_thread.start()
1036
+
1037
+ # Return job ID immediately
1038
+ return jsonify({"job_id": job_id, "check_progress_at": f"/progress/{job_id}"}), 202
1039
 
1040
+ @app.route('/compare-download/<job_id>/<detail_level>', methods=['GET'])
1041
+ def download_comparison_model(job_id, detail_level):
1042
  if job_id not in processing_jobs or processing_jobs[job_id]['status'] != 'completed':
1043
+ return jsonify({"error": "Model not found or processing not complete"}), 404
1044
 
1045
+ if 'comparison' not in processing_jobs[job_id] or not processing_jobs[job_id]['comparison']:
1046
+ return jsonify({"error": "This is not a comparison job"}), 400
1047
+
1048
+ if detail_level not in ['low', 'medium', 'high']:
1049
+ return jsonify({"error": "Invalid detail level"}), 400
1050
+
1051
+ # Get the output directory for this job
1052
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
1053
+ model_path = os.path.join(output_dir, f"model_{detail_level}.glb")
1054
+
1055
+ if os.path.exists(model_path):
1056
+ return send_file(model_path, as_attachment=True, download_name=f"model_{detail_level}.glb")
1057
+
1058
+ return jsonify({"error": "File not found"}), 404
1059
 
1060
  if __name__ == '__main__':
1061
+ # Start the cleanup thread
1062
+ cleanup_old_jobs()
1063
+
1064
+ # Use port 7860 which is standard for Hugging Face Spaces
1065
+ port = int(os.environ.get('PORT', 7860))
1066
+ app.run(host='0.0.0.0', port=port)