update

app.py

@@ -34,20 +34,46 @@ def start_session(req: gr.Request):
 def end_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     shutil.rmtree(user_dir)
+    
+
+@spaces.GPU()
+def remove_background(input: Image.Image) -> Image.Image:
+    input = input.convert('RGB')
+    output = pipeline.rembg_model(input)
+    return output
 
 
-def preprocess_image(image: Image.Image) -> Image.Image:
+def preprocess_image(self, input: Image.Image) -> Image.Image:
     """
     Preprocess the input image.
-
-    Args:
-        image (Image.Image): The input image.
-
-    Returns:
-        Image.Image: The preprocessed image.
     """
-    processed_image = pipeline.preprocess_image(image)
-    return processed_image
+    # if has alpha channel, use it directly; otherwise, remove background
+    has_alpha = False
+    if input.mode == 'RGBA':
+        alpha = np.array(input)[:, :, 3]
+        if not np.all(alpha == 255):
+            has_alpha = True
+    max_size = max(input.size)
+    scale = min(1, 1024 / max_size)
+    if scale < 1:
+        input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
+    if has_alpha:
+        output = input
+    else:
+        output = remove_background(input)
+    output_np = np.array(output)
+    alpha = output_np[:, :, 3]
+    bbox = np.argwhere(alpha > 0.8 * 255)
+    bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
+    center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
+    size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
+    size = int(size * 1)
+    bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
+    output = output.crop(bbox)  # type: ignore
+    output = np.array(output).astype(np.float32) / 255
+    output = output[:, :, :3] * output[:, :, 3:4]
+    output = Image.fromarray((output * 255).astype(np.uint8))
+    return output
 
 
 def pack_state(latents: Tuple[SparseTensor, SparseTensor, int]) -> dict:

trellis2/pipelines/trellis2_image_to_3d.py

@@ -11,13 +11,6 @@ from ..modules import image_feature_extractor
 from ..representations import Mesh, MeshWithVoxel
 
 
-@spaces.GPU()
-def remove_background(rembg_model, input: Image.Image) -> Image.Image:
-    input = input.convert('RGB')
-    output = rembg_model(input)
-    return output
-
-
 class Trellis2ImageTo3DPipeline(Pipeline):
     """
     Pipeline for inferring Trellis2 image-to-3D models.
@@ -139,9 +132,10 @@ class Trellis2ImageTo3DPipeline(Pipeline):
         if has_alpha:
             output = input
         else:
+            input = input.convert('RGB')
             if self.low_vram:
                 self.rembg_model.to(self.device)
-            output = remove_background(self.rembg_model, input)
+            output = self.rembg_model(input)
             if self.low_vram:
                 self.rembg_model.cpu()
         output_np = np.array(output)