Update app_3.py

app_3.py

@@ -780,6 +780,9 @@ def process_bg(input_fg, input_bg, prompt, image_width, image_height, num_sample
 
 @torch.inference_mode()
 def process_relight(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source):
+    # Convert input foreground from PIL to NumPy array if it's in PIL format
+    if isinstance(input_fg, Image.Image):
+        input_fg = np.array(input_fg)
     logging.info(f"Input foreground shape: {input_fg.shape}, dtype: {input_fg.dtype}")
     results = process(input_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source)
     logging.info(f"Results shape: {results.shape}, dtype: {results.dtype}")
@@ -993,6 +996,8 @@ def use_orientation(selected_image:gr.SelectData):
 @torch.inference_mode
 def process_image(input_image, input_text):
     """Main processing function for the Gradio interface"""
+    if isinstance(input_image, Image.Image):
+        input_image = np.array(input_image)
 
     # Initialize configs
     API_TOKEN = "9c8c865e10ec1821bea79d9fa9dc8720"
@@ -1003,7 +1008,6 @@ def process_image(input_image, input_text):
     OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
 
-
     # Initialize DDS client
     config = Config(API_TOKEN)
     client = Client(config)
@@ -1013,6 +1017,8 @@ def process_image(input_image, input_text):
     class_name_to_id = {name: id for id, name in enumerate(classes)}
     class_id_to_name = {id: name for name, id in class_name_to_id.items()}
 
+    
+
     # Save input image to temp file and get URL
     with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmpfile:
         cv2.imwrite(tmpfile.name, input_image)
@@ -1078,12 +1084,41 @@ def process_image(input_image, input_text):
             
             # Get original RGB image
             img = input_image.copy()
-            H, W, C = img.shape
+
+            # rgba = preprocess_image(img)
+
+            alpha = img[..., 3] > 0
+            H, W = alpha.shape
+            # get the bounding box of alpha
+            y, x = np.where(alpha)
+            y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
+            x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
+            image_center = img[y0:y1, x0:x1]
+            # resize the longer side to H * 0.9
+            H, W, _ = image_center.shape
+            if H > W:
+                W = int(W * (height * 0.9) / H)
+                H = int(height * 0.9)
+            else:
+                H = int(H * (width * 0.9) / W)
+                W = int(width * 0.9)
+            image_center = np.array(Image.fromarray(image_center).resize((W, H)))
+            # pad to H, W
+            start_h = (height - H) // 2
+            start_w = (width - W) // 2
+            image = np.zeros((height, width, 4), dtype=np.uint8)
+            image[start_h : start_h + H, start_w : start_w + W] = image_center
+            image = image.astype(np.float32) / 255.0
+            image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
+            image = (image * 255).clip(0, 255).astype(np.uint8)
+            image = Image.fromarray(image)
             
-            # Create RGBA image
-            alpha = np.zeros((H, W, 1), dtype=np.uint8)
-            alpha[first_mask] = 255
-            rgba = np.dstack((img, alpha)).astype(np.uint8)
+            # H, W, C = img.shape
+            
+            # # Create RGBA image
+            # alpha = np.zeros((H, W, 1), dtype=np.uint8)
+            # alpha[first_mask] = 255
+            # rgba = np.dstack((img, alpha)).astype(np.uint8)
             
             # Crop to mask bounds to minimize image size
             # y_indices, x_indices = np.where(first_mask)