Refactor image processing functions in app.py to improve functionality and clarity. Updated preprocess_image and gen_image methods for better handling of background choices and image generation, while enhancing error handling and overall structure.

app.py

@@ -146,89 +146,33 @@ def add_random_background(image, color):
     background = Image.new("RGBA", image.size, color)
     return Image.alpha_composite(background, image)
 
-@spaces.GPU
-def preprocess_image(input_image, background_choice, foreground_ratio, back_ground_color):
-    """Preprocess the input image"""
-    try:
-        # Check if image is provided
-        if input_image is None:
-            raise gr.Error("No image uploaded!")
-            
-        # Convert to PIL Image if needed
-        if isinstance(input_image, np.ndarray):
-            input_image = Image.fromarray(input_image)
-        
-        # Ensure RGBA mode
-        if input_image.mode != "RGBA":
-            input_image = input_image.convert("RGBA")
-        
-        # Process background
-        if background_choice == "Auto Remove background":
-            input_image = remove_background(input_image, rembg_session)
-        elif background_choice == "Custom Background":
-            input_image = add_random_background(input_image, back_ground_color)
-            
-        # Resize content if needed
-        if foreground_ratio != 1.0:
-            input_image = do_resize_content(input_image, foreground_ratio)
-            
-        return input_image
-    except Exception as e:
-        print(f"Error in preprocess_image: {str(e)}")
-        raise gr.Error(f"Preprocessing failed: {str(e)}")
+def preprocess_image(image, background_choice, foreground_ratio, backgroud_color):
+    """
+    input image is a pil image in RGBA, return RGB image
+    """
+    print(background_choice)
+    if background_choice == "Alpha as mask":
+        background = Image.new("RGBA", image.size, (0, 0, 0, 0))
+        image = Image.alpha_composite(background, image)
+    else:
+        image = remove_background(image, rembg_session, force=True)
+    image = do_resize_content(image, foreground_ratio)
+    image = expand_to_square(image)
+    image = add_background(image, backgroud_color)
+    return image.convert("RGB")
 
 @spaces.GPU
-def gen_image(processed_image, seed, scale, step):
-    """Generate the 3D model"""
-    try:
-        # Get model and pipeline when needed
-        model, pipeline = get_model()
-        
-        # Check if image is provided
-        if processed_image is None:
-            raise gr.Error("No processed image provided!")
-        
-        # Convert to numpy array
-        if isinstance(processed_image, Image.Image):
-            np_image = np.array(processed_image)
-        else:
-            np_image = processed_image
-        
-        # Set random seed
-        torch.manual_seed(int(seed))
-        np.random.seed(int(seed))
-        
-        # Generate images
-        np_imgs, np_xyzs = pipeline.generate(
-            np_image,
-            guidance_scale=float(scale),
-            num_inference_steps=int(step)
-        )
-        
-        # Generate 3D model
-        glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
-        
-        return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path
-    except Exception as e:
-        print(f"Error in gen_image: {str(e)}")
-        raise gr.Error(f"Generation failed: {str(e)}")
+def gen_image(input_image, seed, scale, step):
+    global pipeline, model, args
+    pipeline.set_seed(seed)
+    rt_dict = pipeline(input_image, scale=scale, step=step)
+    stage1_images = rt_dict["stage1_images"]
+    stage2_images = rt_dict["stage2_images"]
+    np_imgs = np.concatenate(stage1_images, 1)
+    np_xyzs = np.concatenate(stage2_images, 1)
 
-def process_and_generate(image, bg_choice, fg_ratio, bg_color, seed_val, guidance, steps):
-    """Combined function to process image and generate 3D model"""
-    try:
-        if image is None:
-            raise gr.Error("No image uploaded!")
-        
-        # Preprocess the image
-        processed = preprocess_image(image, bg_choice, fg_ratio, bg_color)
-        
-        # Generate 3D model
-        rgb_img, ccm_img, glb_file = gen_image(processed, seed_val, guidance, steps)
-        
-        return processed, rgb_img, ccm_img, glb_file
-    except Exception as e:
-        print(f"Error in process_and_generate: {str(e)}")
-        raise gr.Error(f"Process failed: {str(e)}")
+    glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
+    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path
 
 _DESCRIPTION = '''
 * Our [official implementation](https://github.com/thu-ml/CRM) uses UV texture instead of vertex color. It has better texture than this online demo.