mac9087 commited on
Commit
1c91a49
·
verified ·
1 Parent(s): 7798d49

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +430 -62
app.py CHANGED
@@ -16,7 +16,9 @@ from flask_cors import CORS
16
  import numpy as np
17
  import trimesh
18
  from transformers import pipeline
19
- from scipy.ndimage import gaussian_filter, uniform_filter
 
 
20
 
21
  app = Flask(__name__)
22
  CORS(app) # Enable CORS for all routes
@@ -49,8 +51,8 @@ model_loaded = False
49
  model_loading = False
50
 
51
  # Configuration for processing
52
- TIMEOUT_SECONDS = 240 # 4 minutes max for processing (increased for larger model)
53
- MAX_DIMENSION = 512 # Max image dimension to processe
54
 
55
  # TimeoutError for handling timeouts
56
  class TimeoutError(Exception):
@@ -89,10 +91,11 @@ def process_with_timeout(function, args, timeout):
89
  def allowed_file(filename):
90
  return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
91
 
92
- # Function to preprocess image
93
  def preprocess_image(image_path):
94
  with Image.open(image_path) as img:
95
  img = img.convert("RGB")
 
96
  # Resize if the image is too large
97
  if img.width > MAX_DIMENSION or img.height > MAX_DIMENSION:
98
  # Calculate new dimensions while preserving aspect ratio
@@ -102,8 +105,33 @@ def preprocess_image(image_path):
102
  else:
103
  new_height = MAX_DIMENSION
104
  new_width = int(img.width * (MAX_DIMENSION / img.height))
 
 
105
  img = img.resize((new_width, new_height), Image.LANCZOS)
106
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
107
  return img
108
 
109
  def load_model():
@@ -123,6 +151,7 @@ def load_model():
123
  print("Starting model loading...")
124
 
125
  # Using DPT-Large which provides better detail than DPT-Hybrid
 
126
  model_name = "Intel/dpt-large"
127
 
128
  # Download model with retry mechanism
@@ -145,7 +174,7 @@ def load_model():
145
  else:
146
  raise
147
 
148
- # Initialize model with lower precision to save memory
149
  device = "cuda" if torch.cuda.is_available() else "cpu"
150
 
151
  # Load depth estimator pipeline
@@ -171,10 +200,10 @@ def load_model():
171
  finally:
172
  model_loading = False
173
 
174
- # Convert depth map to 3D mesh with enhanced detail
175
- def depth_to_mesh(depth_map, image, resolution=100):
176
- """Convert depth map to 3D mesh with improved detail preservation"""
177
- # Convert depth_map to numpy array if it's a PIL Image
178
  if isinstance(depth_map, Image.Image):
179
  depth_map = np.array(depth_map)
180
 
@@ -182,39 +211,101 @@ def depth_to_mesh(depth_map, image, resolution=100):
182
  if len(depth_map.shape) > 2:
183
  depth_map = np.mean(depth_map, axis=2) if depth_map.shape[2] > 1 else depth_map[:,:,0]
184
 
185
- # Apply bilateral filter to smooth the depth map while preserving edges
186
- # First, apply a slight gaussian filter to remove noise
187
- depth_map_smooth = gaussian_filter(depth_map, sigma=1.0)
 
 
 
 
 
 
188
 
189
- # Get dimensions
190
- h, w = depth_map_smooth.shape
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
191
 
192
- # Create a grid of points
 
 
 
 
 
 
 
 
 
 
 
 
 
 
193
  x = np.linspace(0, w-1, resolution)
194
  y = np.linspace(0, h-1, resolution)
195
  x_grid, y_grid = np.meshgrid(x, y)
196
 
197
- # Sample depth at grid points
198
- x_indices = x_grid.astype(int)
199
- y_indices = y_grid.astype(int)
200
- z_values = depth_map_smooth[y_indices, x_indices]
 
 
 
 
201
 
202
- # Normalize depth values with better scaling
203
- z_min, z_max = np.percentile(z_values, [2, 98]) # Removes outliers
 
 
 
 
 
 
 
 
 
 
 
 
 
204
  z_values = (z_values - z_min) / (z_max - z_min) if z_max > z_min else z_values
205
- z_values = z_values * 2.0 # Scale depth
206
-
207
- # Apply a local contrast enhancement to bring out details
208
- # Simple adaptive normalization
209
- window_size = resolution // 10
210
- if window_size > 0:
211
- local_mean = uniform_filter(z_values, size=window_size)
212
- local_var = uniform_filter(z_values**2, size=window_size) - local_mean**2
213
- local_std = np.sqrt(np.maximum(local_var, 0))
214
 
215
- # Enhance local contrast
216
- enhanced_z = (z_values - local_mean) / (local_std + 0.01) * 0.5 + z_values
217
- z_values = np.clip(enhanced_z, 0, None) # Keep values positive
218
 
219
  # Normalize x and y coordinates
220
  x_grid = (x_grid / w - 0.5) * 2.0 # Map to -1 to 1
@@ -223,7 +314,7 @@ def depth_to_mesh(depth_map, image, resolution=100):
223
  # Create vertices
224
  vertices = np.vstack([x_grid.flatten(), -y_grid.flatten(), -z_values.flatten()]).T
225
 
226
- # Create faces (triangles)
227
  faces = []
228
  for i in range(resolution-1):
229
  for j in range(resolution-1):
@@ -232,54 +323,103 @@ def depth_to_mesh(depth_map, image, resolution=100):
232
  p3 = (i + 1) * resolution + j
233
  p4 = (i + 1) * resolution + (j + 1)
234
 
235
- # Create two triangles for each grid cell
236
- faces.append([p1, p2, p4])
237
- faces.append([p1, p4, p3])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
238
 
239
  faces = np.array(faces)
240
 
241
  # Create mesh
242
  mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
243
 
244
- # Apply texturing if image is provided
245
  if image:
246
  # Convert to numpy array if needed
247
  if isinstance(image, Image.Image):
248
  img_array = np.array(image)
249
  else:
250
  img_array = image
251
-
252
- # Create simple texture by sampling the original image
253
  if resolution <= img_array.shape[0] and resolution <= img_array.shape[1]:
254
- # Create vertex colors by sampling the image
255
  vertex_colors = np.zeros((vertices.shape[0], 4), dtype=np.uint8)
256
 
 
257
  for i in range(resolution):
258
  for j in range(resolution):
259
- img_x = min(int(j * img_array.shape[1] / resolution), img_array.shape[1]-1)
260
- img_y = min(int(i * img_array.shape[0] / resolution), img_array.shape[0]-1)
 
 
 
 
 
 
 
 
 
261
 
262
  vertex_idx = i * resolution + j
 
263
  if len(img_array.shape) == 3 and img_array.shape[2] == 3: # RGB
264
- vertex_colors[vertex_idx, :3] = img_array[img_y, img_x, :]
 
 
 
 
 
 
 
 
265
  vertex_colors[vertex_idx, 3] = 255 # Alpha
266
  elif len(img_array.shape) == 3 and img_array.shape[2] == 4: # RGBA
267
- vertex_colors[vertex_idx, :] = img_array[img_y, img_x, :]
 
 
 
 
268
  else:
269
- # Handle grayscale or other formats
270
- gray_value = img_array[img_y, img_x]
271
- vertex_colors[vertex_idx, :3] = [gray_value, gray_value, gray_value]
 
272
  vertex_colors[vertex_idx, 3] = 255
273
 
274
  mesh.visual.vertex_colors = vertex_colors
275
 
 
 
 
 
 
 
 
 
 
276
  return mesh
277
 
278
  @app.route('/health', methods=['GET'])
279
  def health_check():
280
  return jsonify({
281
  "status": "healthy",
282
- "model": "Depth-Based 3D Model Generator (DPT-Large)",
283
  "device": "cuda" if torch.cuda.is_available() else "cpu"
284
  }), 200
285
 
@@ -339,7 +479,8 @@ def convert_image_to_3d():
339
  try:
340
  mesh_resolution = min(int(request.form.get('mesh_resolution', 100)), 200) # Limit max resolution
341
  output_format = request.form.get('output_format', 'obj').lower()
342
- detail_level = request.form.get('detail_level', 'medium').lower() # New parameter for detail level
 
343
  except ValueError:
344
  return jsonify({"error": "Invalid parameter values"}), 400
345
 
@@ -349,7 +490,7 @@ def convert_image_to_3d():
349
 
350
  # Adjust mesh resolution based on detail level
351
  if detail_level == 'high':
352
- mesh_resolution = min(mesh_resolution * 1.5, 200)
353
  elif detail_level == 'low':
354
  mesh_resolution = max(int(mesh_resolution * 0.7), 50)
355
 
@@ -380,7 +521,7 @@ def convert_image_to_3d():
380
  processing_jobs[job_id]['thread_alive'] = lambda: thread.is_alive()
381
 
382
  try:
383
- # Preprocess image
384
  processing_jobs[job_id]['progress'] = 5
385
  image = preprocess_image(filepath)
386
  processing_jobs[job_id]['progress'] = 10
@@ -423,9 +564,9 @@ def convert_image_to_3d():
423
 
424
  processing_jobs[job_id]['progress'] = 60
425
 
426
- # Create mesh from depth map
427
  mesh_resolution_int = int(mesh_resolution)
428
- mesh = depth_to_mesh(depth_map, image, resolution=mesh_resolution_int)
429
  processing_jobs[job_id]['progress'] = 80
430
 
431
  except Exception as e:
@@ -436,11 +577,18 @@ def convert_image_to_3d():
436
  print(error_details)
437
  return
438
 
439
- # Export based on requested format
440
  try:
441
  if output_format == 'obj':
442
  obj_path = os.path.join(output_dir, "model.obj")
443
- mesh.export(obj_path, file_type='obj')
 
 
 
 
 
 
 
444
 
445
  # Create a zip file with OBJ and MTL
446
  zip_path = os.path.join(output_dir, "model.zip")
@@ -449,14 +597,22 @@ def convert_image_to_3d():
449
  mtl_path = os.path.join(output_dir, "model.mtl")
450
  if os.path.exists(mtl_path):
451
  zipf.write(mtl_path, arcname="model.mtl")
 
 
 
 
 
452
 
453
  processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
454
  processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
455
 
456
  elif output_format == 'glb':
457
- # Export as GLB
458
  glb_path = os.path.join(output_dir, "model.glb")
459
- mesh.export(glb_path, file_type='glb')
 
 
 
460
 
461
  processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
462
  processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
@@ -573,18 +729,230 @@ def cleanup_old_jobs():
573
  # Schedule the next cleanup
574
  threading.Timer(300, cleanup_old_jobs).start() # Run every 5 minutes
575
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
576
  @app.route('/', methods=['GET'])
577
  def index():
578
  return jsonify({
579
- "message": "Image to 3D API is running (DPT-Large Model)",
580
- "endpoints": ["/convert", "/progress/<job_id>", "/download/<job_id>", "/preview/<job_id>"],
 
 
 
 
 
 
581
  "parameters": {
582
  "mesh_resolution": "Integer (50-200), controls mesh density",
583
  "output_format": "obj or glb",
584
- "detail_level": "low, medium, or high"
585
- }
 
 
586
  }), 200
587
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
588
  if __name__ == '__main__':
589
  # Start the cleanup thread
590
  cleanup_old_jobs()
 
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
 
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):
 
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
  img = img.convert("RGB")
98
+
99
  # Resize if the image is too large
100
  if img.width > MAX_DIMENSION or img.height > MAX_DIMENSION:
101
  # Calculate new dimensions while preserving aspect ratio
 
105
  else:
106
  new_height = MAX_DIMENSION
107
  new_width = int(img.width * (MAX_DIMENSION / img.height))
108
+
109
+ # Use high-quality Lanczos resampling for better detail preservation
110
  img = img.resize((new_width, new_height), Image.LANCZOS)
111
 
112
+ # Convert to numpy array for additional preprocessing
113
+ img_array = np.array(img)
114
+
115
+ # Optional: Apply adaptive histogram equalization for better contrast
116
+ # This helps the depth model detect more details
117
+ if len(img_array.shape) == 3 and img_array.shape[2] == 3:
118
+ # Convert to LAB color space
119
+ lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
120
+ l, a, b = cv2.split(lab)
121
+
122
+ # Apply CLAHE to L channel
123
+ clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
124
+ cl = clahe.apply(l)
125
+
126
+ # Merge channels back
127
+ enhanced_lab = cv2.merge((cl, a, b))
128
+
129
+ # Convert back to RGB
130
+ img_array = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2RGB)
131
+
132
+ # Convert back to PIL Image
133
+ img = Image.fromarray(img_array)
134
+
135
  return img
136
 
137
  def load_model():
 
151
  print("Starting model loading...")
152
 
153
  # Using DPT-Large which provides better detail than DPT-Hybrid
154
+ # Alternatively, consider "vinvino02/glpn-nyu" for different detail characteristics
155
  model_name = "Intel/dpt-large"
156
 
157
  # Download model with retry mechanism
 
174
  else:
175
  raise
176
 
177
+ # Initialize model with appropriate precision
178
  device = "cuda" if torch.cuda.is_available() else "cpu"
179
 
180
  # Load depth estimator pipeline
 
200
  finally:
201
  model_loading = False
202
 
203
+ # Enhanced depth processing function to improve detail quality
204
+ def enhance_depth_map(depth_map, detail_level='medium'):
205
+ """Apply sophisticated processing to enhance depth map details"""
206
+ # Convert to numpy array if needed
207
  if isinstance(depth_map, Image.Image):
208
  depth_map = np.array(depth_map)
209
 
 
211
  if len(depth_map.shape) > 2:
212
  depth_map = np.mean(depth_map, axis=2) if depth_map.shape[2] > 1 else depth_map[:,:,0]
213
 
214
+ # Create a copy for processing
215
+ enhanced_depth = depth_map.copy().astype(np.float32)
216
+
217
+ # Remove outliers using percentile clipping (more stable than min/max)
218
+ p_low, p_high = np.percentile(enhanced_depth, [1, 99])
219
+ enhanced_depth = np.clip(enhanced_depth, p_low, p_high)
220
+
221
+ # Normalize to 0-1 range for processing
222
+ enhanced_depth = (enhanced_depth - p_low) / (p_high - p_low) if p_high > p_low else enhanced_depth
223
 
224
+ # Apply different enhancement methods based on detail level
225
+ if detail_level == 'high':
226
+ # Apply unsharp masking for edge enhancement - simulating Hunyuan's detail technique
227
+ # First apply gaussian blur
228
+ blurred = gaussian_filter(enhanced_depth, sigma=1.5)
229
+ # Create the unsharp mask
230
+ mask = enhanced_depth - blurred
231
+ # Apply the mask with strength factor
232
+ enhanced_depth = enhanced_depth + 1.5 * mask
233
+
234
+ # Apply bilateral filter to preserve edges while smoothing noise
235
+ # Simulate using gaussian combinations
236
+ smooth1 = gaussian_filter(enhanced_depth, sigma=0.5)
237
+ smooth2 = gaussian_filter(enhanced_depth, sigma=2.0)
238
+ edge_mask = enhanced_depth - smooth2
239
+ enhanced_depth = smooth1 + 1.2 * edge_mask
240
+
241
+ elif detail_level == 'medium':
242
+ # Less aggressive but still effective enhancement
243
+ # Apply mild unsharp masking
244
+ blurred = gaussian_filter(enhanced_depth, sigma=1.0)
245
+ mask = enhanced_depth - blurred
246
+ enhanced_depth = enhanced_depth + 0.8 * mask
247
+
248
+ # Apply mild smoothing to reduce noise but preserve edges
249
+ enhanced_depth = gaussian_filter(enhanced_depth, sigma=0.5)
250
+
251
+ else: # low
252
+ # Just apply noise reduction without too much detail enhancement
253
+ enhanced_depth = gaussian_filter(enhanced_depth, sigma=0.7)
254
 
255
+ # Normalize again after processing
256
+ enhanced_depth = np.clip(enhanced_depth, 0, 1)
257
+
258
+ return enhanced_depth
259
+
260
+ # Convert depth map to 3D mesh with significantly enhanced detail
261
+ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
262
+ """Convert depth map to 3D mesh with highly improved detail preservation"""
263
+ # First, enhance the depth map for better details
264
+ enhanced_depth = enhance_depth_map(depth_map, detail_level)
265
+
266
+ # Get dimensions of depth map
267
+ h, w = enhanced_depth.shape
268
+
269
+ # Create a higher resolution grid for better detail
270
  x = np.linspace(0, w-1, resolution)
271
  y = np.linspace(0, h-1, resolution)
272
  x_grid, y_grid = np.meshgrid(x, y)
273
 
274
+ # Use bicubic interpolation for smoother surface with better details
275
+ # Create interpolation function
276
+ interp_func = interpolate.RectBivariateSpline(
277
+ np.arange(h), np.arange(w), enhanced_depth, kx=3, ky=3
278
+ )
279
+
280
+ # Sample depth at grid points with the interpolation function
281
+ z_values = interp_func(y, x, grid=True)
282
 
283
+ # Apply a post-processing step to enhance small details even further
284
+ if detail_level == 'high':
285
+ # Calculate local gradients to detect edges
286
+ dx = np.gradient(z_values, axis=1)
287
+ dy = np.gradient(z_values, axis=0)
288
+
289
+ # Enhance edges by increasing depth differences at high gradient areas
290
+ gradient_magnitude = np.sqrt(dx**2 + dy**2)
291
+ edge_mask = np.clip(gradient_magnitude * 5, 0, 0.2) # Scale and limit effect
292
+
293
+ # Apply edge enhancement
294
+ z_values = z_values + edge_mask * (z_values - gaussian_filter(z_values, sigma=1.0))
295
+
296
+ # Normalize z-values with advanced scaling for better depth impression
297
+ z_min, z_max = np.percentile(z_values, [2, 98]) # Remove outliers
298
  z_values = (z_values - z_min) / (z_max - z_min) if z_max > z_min else z_values
299
+
300
+ # Apply depth scaling appropriate to the detail level
301
+ if detail_level == 'high':
302
+ z_scaling = 2.5 # More pronounced depth variations
303
+ elif detail_level == 'medium':
304
+ z_scaling = 2.0 # Standard depth
305
+ else:
306
+ z_scaling = 1.5 # More subtle depth variations
 
307
 
308
+ z_values = z_values * z_scaling
 
 
309
 
310
  # Normalize x and y coordinates
311
  x_grid = (x_grid / w - 0.5) * 2.0 # Map to -1 to 1
 
314
  # Create vertices
315
  vertices = np.vstack([x_grid.flatten(), -y_grid.flatten(), -z_values.flatten()]).T
316
 
317
+ # Create faces (triangles) with optimized winding for better normals
318
  faces = []
319
  for i in range(resolution-1):
320
  for j in range(resolution-1):
 
323
  p3 = (i + 1) * resolution + j
324
  p4 = (i + 1) * resolution + (j + 1)
325
 
326
+ # Calculate normals to ensure consistent orientation
327
+ v1 = vertices[p1]
328
+ v2 = vertices[p2]
329
+ v3 = vertices[p3]
330
+ v4 = vertices[p4]
331
+
332
+ # Calculate normals for both possible triangulations
333
+ # and choose the one that's more consistent
334
+ norm1 = np.cross(v2-v1, v4-v1)
335
+ norm2 = np.cross(v4-v3, v1-v3)
336
+
337
+ if np.dot(norm1, norm2) >= 0:
338
+ # Standard triangulation
339
+ faces.append([p1, p2, p4])
340
+ faces.append([p1, p4, p3])
341
+ else:
342
+ # Alternative triangulation for smoother surface
343
+ faces.append([p1, p2, p3])
344
+ faces.append([p2, p4, p3])
345
 
346
  faces = np.array(faces)
347
 
348
  # Create mesh
349
  mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
350
 
351
+ # Apply advanced texturing if image is provided
352
  if image:
353
  # Convert to numpy array if needed
354
  if isinstance(image, Image.Image):
355
  img_array = np.array(image)
356
  else:
357
  img_array = image
358
+
359
+ # Create vertex colors with improved sampling
360
  if resolution <= img_array.shape[0] and resolution <= img_array.shape[1]:
361
+ # Create vertex colors by sampling the image with bilinear interpolation
362
  vertex_colors = np.zeros((vertices.shape[0], 4), dtype=np.uint8)
363
 
364
+ # Get normalized coordinates for sampling
365
  for i in range(resolution):
366
  for j in range(resolution):
367
+ # Calculate exact image coordinates with proper scaling
368
+ img_x = j * (img_array.shape[1] - 1) / (resolution - 1)
369
+ img_y = i * (img_array.shape[0] - 1) / (resolution - 1)
370
+
371
+ # Bilinear interpolation for smooth color transitions
372
+ x0, y0 = int(img_x), int(img_y)
373
+ x1, y1 = min(x0 + 1, img_array.shape[1] - 1), min(y0 + 1, img_array.shape[0] - 1)
374
+
375
+ # Calculate interpolation weights
376
+ wx = img_x - x0
377
+ wy = img_y - y0
378
 
379
  vertex_idx = i * resolution + j
380
+
381
  if len(img_array.shape) == 3 and img_array.shape[2] == 3: # RGB
382
+ # Perform bilinear interpolation for each color channel
383
+ r = int((1-wx)*(1-wy)*img_array[y0, x0, 0] + wx*(1-wy)*img_array[y0, x1, 0] +
384
+ (1-wx)*wy*img_array[y1, x0, 0] + wx*wy*img_array[y1, x1, 0])
385
+ g = int((1-wx)*(1-wy)*img_array[y0, x0, 1] + wx*(1-wy)*img_array[y0, x1, 1] +
386
+ (1-wx)*wy*img_array[y1, x0, 1] + wx*wy*img_array[y1, x1, 1])
387
+ b = int((1-wx)*(1-wy)*img_array[y0, x0, 2] + wx*(1-wy)*img_array[y0, x1, 2] +
388
+ (1-wx)*wy*img_array[y1, x0, 2] + wx*wy*img_array[y1, x1, 2])
389
+
390
+ vertex_colors[vertex_idx, :3] = [r, g, b]
391
  vertex_colors[vertex_idx, 3] = 255 # Alpha
392
  elif len(img_array.shape) == 3 and img_array.shape[2] == 4: # RGBA
393
+ for c in range(4): # For each RGBA channel
394
+ vertex_colors[vertex_idx, c] = int((1-wx)*(1-wy)*img_array[y0, x0, c] +
395
+ wx*(1-wy)*img_array[y0, x1, c] +
396
+ (1-wx)*wy*img_array[y1, x0, c] +
397
+ wx*wy*img_array[y1, x1, c])
398
  else:
399
+ # Handle grayscale with bilinear interpolation
400
+ gray = int((1-wx)*(1-wy)*img_array[y0, x0] + wx*(1-wy)*img_array[y0, x1] +
401
+ (1-wx)*wy*img_array[y1, x0] + wx*wy*img_array[y1, x1])
402
+ vertex_colors[vertex_idx, :3] = [gray, gray, gray]
403
  vertex_colors[vertex_idx, 3] = 255
404
 
405
  mesh.visual.vertex_colors = vertex_colors
406
 
407
+ # Apply smoothing to get rid of staircase artifacts
408
+ if detail_level != 'high':
409
+ # For medium and low detail, apply Laplacian smoothing
410
+ # but preserve the overall shape
411
+ mesh = mesh.smoothed(method='laplacian', iterations=1)
412
+
413
+ # Calculate and fix normals for better rendering
414
+ mesh.fix_normals()
415
+
416
  return mesh
417
 
418
  @app.route('/health', methods=['GET'])
419
  def health_check():
420
  return jsonify({
421
  "status": "healthy",
422
+ "model": "Enhanced Depth-Based 3D Model Generator (DPT-Large)",
423
  "device": "cuda" if torch.cuda.is_available() else "cpu"
424
  }), 200
425
 
 
479
  try:
480
  mesh_resolution = min(int(request.form.get('mesh_resolution', 100)), 200) # Limit max resolution
481
  output_format = request.form.get('output_format', 'obj').lower()
482
+ detail_level = request.form.get('detail_level', 'medium').lower() # Parameter for detail level
483
+ texture_quality = request.form.get('texture_quality', 'medium').lower() # New parameter for texture quality
484
  except ValueError:
485
  return jsonify({"error": "Invalid parameter values"}), 400
486
 
 
490
 
491
  # Adjust mesh resolution based on detail level
492
  if detail_level == 'high':
493
+ mesh_resolution = min(int(mesh_resolution * 1.5), 200)
494
  elif detail_level == 'low':
495
  mesh_resolution = max(int(mesh_resolution * 0.7), 50)
496
 
 
521
  processing_jobs[job_id]['thread_alive'] = lambda: thread.is_alive()
522
 
523
  try:
524
+ # Preprocess image with enhanced detail preservation
525
  processing_jobs[job_id]['progress'] = 5
526
  image = preprocess_image(filepath)
527
  processing_jobs[job_id]['progress'] = 10
 
564
 
565
  processing_jobs[job_id]['progress'] = 60
566
 
567
+ # Create mesh from depth map with enhanced detail handling
568
  mesh_resolution_int = int(mesh_resolution)
569
+ mesh = depth_to_mesh(depth_map, image, resolution=mesh_resolution_int, detail_level=detail_level)
570
  processing_jobs[job_id]['progress'] = 80
571
 
572
  except Exception as e:
 
577
  print(error_details)
578
  return
579
 
580
+ # Export based on requested format with enhanced quality settings
581
  try:
582
  if output_format == 'obj':
583
  obj_path = os.path.join(output_dir, "model.obj")
584
+
585
+ # Export with normal and texture coordinates
586
+ mesh.export(
587
+ obj_path,
588
+ file_type='obj',
589
+ include_normals=True,
590
+ include_texture=True
591
+ )
592
 
593
  # Create a zip file with OBJ and MTL
594
  zip_path = os.path.join(output_dir, "model.zip")
 
597
  mtl_path = os.path.join(output_dir, "model.mtl")
598
  if os.path.exists(mtl_path):
599
  zipf.write(mtl_path, arcname="model.mtl")
600
+
601
+ # Include texture file if it exists
602
+ texture_path = os.path.join(output_dir, "model.png")
603
+ if os.path.exists(texture_path):
604
+ zipf.write(texture_path, arcname="model.png")
605
 
606
  processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
607
  processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
608
 
609
  elif output_format == 'glb':
610
+ # Export as GLB with enhanced settings
611
  glb_path = os.path.join(output_dir, "model.glb")
612
+ mesh.export(
613
+ glb_path,
614
+ file_type='glb'
615
+ )
616
 
617
  processing_jobs[job_id]['result_url'] = f"/download/{job_id}"
618
  processing_jobs[job_id]['preview_url'] = f"/preview/{job_id}"
 
729
  # Schedule the next cleanup
730
  threading.Timer(300, cleanup_old_jobs).start() # Run every 5 minutes
731
 
732
+ # New endpoint to get detailed information about a model
733
+ @app.route('/model-info/<job_id>', methods=['GET'])
734
+ def model_info(job_id):
735
+ if job_id not in processing_jobs:
736
+ return jsonify({"error": "Model not found"}), 404
737
+
738
+ job = processing_jobs[job_id]
739
+
740
+ if job['status'] != 'completed':
741
+ return jsonify({
742
+ "status": job['status'],
743
+ "progress": job['progress'],
744
+ "error": job.get('error')
745
+ }), 200
746
+
747
+ # For completed jobs, include information about the model
748
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
749
+ model_stats = {}
750
+
751
+ # Get file size
752
+ if job['output_format'] == 'obj':
753
+ obj_path = os.path.join(output_dir, "model.obj")
754
+ zip_path = os.path.join(output_dir, "model.zip")
755
+
756
+ if os.path.exists(obj_path):
757
+ model_stats['obj_size'] = os.path.getsize(obj_path)
758
+
759
+ if os.path.exists(zip_path):
760
+ model_stats['package_size'] = os.path.getsize(zip_path)
761
+
762
+ else: # glb
763
+ glb_path = os.path.join(output_dir, "model.glb")
764
+ if os.path.exists(glb_path):
765
+ model_stats['model_size'] = os.path.getsize(glb_path)
766
+
767
+ # Return detailed info
768
+ return jsonify({
769
+ "status": job['status'],
770
+ "model_format": job['output_format'],
771
+ "download_url": job['result_url'],
772
+ "preview_url": job['preview_url'],
773
+ "model_stats": model_stats,
774
+ "created_at": job.get('created_at'),
775
+ "completed_at": job.get('completed_at')
776
+ }), 200
777
+
778
  @app.route('/', methods=['GET'])
779
  def index():
780
  return jsonify({
781
+ "message": "Enhanced Image to 3D API (DPT-Large Model)",
782
+ "endpoints": [
783
+ "/convert",
784
+ "/progress/<job_id>",
785
+ "/download/<job_id>",
786
+ "/preview/<job_id>",
787
+ "/model-info/<job_id>"
788
+ ],
789
  "parameters": {
790
  "mesh_resolution": "Integer (50-200), controls mesh density",
791
  "output_format": "obj or glb",
792
+ "detail_level": "low, medium, or high - controls the level of detail in the final model",
793
+ "texture_quality": "low, medium, or high - controls the quality of textures"
794
+ },
795
+ "description": "This API creates high-quality 3D models from 2D images with enhanced detail finishing similar to Hunyuan model"
796
  }), 200
797
 
798
+ # Example endpoint showing how to compare different detail levels
799
+ @app.route('/detail-comparison', methods=['POST'])
800
+ def compare_detail_levels():
801
+ # Check if image is in the request
802
+ if 'image' not in request.files:
803
+ return jsonify({"error": "No image provided"}), 400
804
+
805
+ file = request.files['image']
806
+ if file.filename == '':
807
+ return jsonify({"error": "No image selected"}), 400
808
+
809
+ if not allowed_file(file.filename):
810
+ return jsonify({"error": f"File type not allowed. Supported types: {', '.join(ALLOWED_EXTENSIONS)}"}), 400
811
+
812
+ # Create a job ID
813
+ job_id = str(uuid.uuid4())
814
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
815
+ os.makedirs(output_dir, exist_ok=True)
816
+
817
+ # Save the uploaded file
818
+ filename = secure_filename(file.filename)
819
+ filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{job_id}_{filename}")
820
+ file.save(filepath)
821
+
822
+ # Initialize job tracking
823
+ processing_jobs[job_id] = {
824
+ 'status': 'processing',
825
+ 'progress': 0,
826
+ 'result_url': None,
827
+ 'preview_url': None,
828
+ 'error': None,
829
+ 'output_format': 'glb', # Use GLB for comparison
830
+ 'created_at': time.time(),
831
+ 'comparison': True
832
+ }
833
+
834
+ # Process in separate thread to create 3 different detail levels
835
+ def process_comparison():
836
+ thread = threading.current_thread()
837
+ processing_jobs[job_id]['thread_alive'] = lambda: thread.is_alive()
838
+
839
+ try:
840
+ # Preprocess image
841
+ image = preprocess_image(filepath)
842
+ processing_jobs[job_id]['progress'] = 10
843
+
844
+ # Load model
845
+ try:
846
+ model = load_model()
847
+ processing_jobs[job_id]['progress'] = 20
848
+ except Exception as e:
849
+ processing_jobs[job_id]['status'] = 'error'
850
+ processing_jobs[job_id]['error'] = f"Error loading model: {str(e)}"
851
+ return
852
+
853
+ # Process image to get depth map
854
+ try:
855
+ depth_map = model(image)["depth"]
856
+ if isinstance(depth_map, torch.Tensor):
857
+ depth_map = depth_map.cpu().numpy()
858
+ elif hasattr(depth_map, 'numpy'):
859
+ depth_map = depth_map.numpy()
860
+ elif isinstance(depth_map, Image.Image):
861
+ depth_map = np.array(depth_map)
862
+
863
+ processing_jobs[job_id]['progress'] = 40
864
+ except Exception as e:
865
+ processing_jobs[job_id]['status'] = 'error'
866
+ processing_jobs[job_id]['error'] = f"Error estimating depth: {str(e)}"
867
+ return
868
+
869
+ # Create meshes at different detail levels
870
+ result_urls = {}
871
+
872
+ for detail_level in ['low', 'medium', 'high']:
873
+ try:
874
+ # Update progress
875
+ if detail_level == 'low':
876
+ processing_jobs[job_id]['progress'] = 50
877
+ elif detail_level == 'medium':
878
+ processing_jobs[job_id]['progress'] = 70
879
+ else:
880
+ processing_jobs[job_id]['progress'] = 90
881
+
882
+ # Create mesh with appropriate detail level
883
+ mesh_resolution = 100 # Fixed resolution for fair comparison
884
+ if detail_level == 'high':
885
+ mesh_resolution = 150
886
+ elif detail_level == 'low':
887
+ mesh_resolution = 80
888
+
889
+ mesh = depth_to_mesh(depth_map, image,
890
+ resolution=mesh_resolution,
891
+ detail_level=detail_level)
892
+
893
+ # Export as GLB
894
+ model_path = os.path.join(output_dir, f"model_{detail_level}.glb")
895
+ mesh.export(model_path, file_type='glb')
896
+
897
+ # Add to result URLs
898
+ result_urls[detail_level] = f"/compare-download/{job_id}/{detail_level}"
899
+
900
+ except Exception as e:
901
+ print(f"Error processing {detail_level} detail level: {str(e)}")
902
+ # Continue with other detail levels even if one fails
903
+
904
+ # Update job status
905
+ processing_jobs[job_id]['status'] = 'completed'
906
+ processing_jobs[job_id]['progress'] = 100
907
+ processing_jobs[job_id]['result_urls'] = result_urls
908
+ processing_jobs[job_id]['completed_at'] = time.time()
909
+
910
+ # Clean up temporary file
911
+ if os.path.exists(filepath):
912
+ os.remove(filepath)
913
+
914
+ # Force garbage collection
915
+ gc.collect()
916
+ if torch.cuda.is_available():
917
+ torch.cuda.empty_cache()
918
+
919
+ except Exception as e:
920
+ # Handle errors
921
+ processing_jobs[job_id]['status'] = 'error'
922
+ processing_jobs[job_id]['error'] = f"Error during processing: {str(e)}"
923
+
924
+ # Clean up on error
925
+ if os.path.exists(filepath):
926
+ os.remove(filepath)
927
+
928
+ # Start processing thread
929
+ processing_thread = threading.Thread(target=process_comparison)
930
+ processing_thread.daemon = True
931
+ processing_thread.start()
932
+
933
+ # Return job ID immediately
934
+ return jsonify({"job_id": job_id, "check_progress_at": f"/progress/{job_id}"}), 202
935
+
936
+ @app.route('/compare-download/<job_id>/<detail_level>', methods=['GET'])
937
+ def download_comparison_model(job_id, detail_level):
938
+ if job_id not in processing_jobs or processing_jobs[job_id]['status'] != 'completed':
939
+ return jsonify({"error": "Model not found or processing not complete"}), 404
940
+
941
+ if 'comparison' not in processing_jobs[job_id] or not processing_jobs[job_id]['comparison']:
942
+ return jsonify({"error": "This is not a comparison job"}), 400
943
+
944
+ if detail_level not in ['low', 'medium', 'high']:
945
+ return jsonify({"error": "Invalid detail level"}), 400
946
+
947
+ # Get the output directory for this job
948
+ output_dir = os.path.join(RESULTS_FOLDER, job_id)
949
+ model_path = os.path.join(output_dir, f"model_{detail_level}.glb")
950
+
951
+ if os.path.exists(model_path):
952
+ return send_file(model_path, as_attachment=True, download_name=f"model_{detail_level}.glb")
953
+
954
+ return jsonify({"error": "File not found"}), 404
955
+
956
  if __name__ == '__main__':
957
  # Start the cleanup thread
958
  cleanup_old_jobs()