Update app_3.py

app_3.py

@@ -542,7 +542,6 @@ def preprocess_image(image: Image.Image, height=768, width=768):
     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)
-
     return image
   
 
@@ -996,6 +995,7 @@ 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)
 
@@ -1085,14 +1085,34 @@ def process_image(input_image, input_text):
             # Get original RGB image
             img = input_image.copy()
 
-            # rgba = preprocess_image(img)
 
-            alpha = img[..., 3] > 0
-            H, W = alpha.shape
+            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)
+            # y_min, y_max = y_indices.min(), y_indices.max()
+            # x_min, x_max = x_indices.min(), x_indices.max()
+            
+            # Crop the RGBA image
+            # cropped_rgba = rgba[y_min:y_max+1, x_min:x_max+1]
+            
+            # Set extracted foreground for mask mover
+            # mask_mover.set_extracted_fg(cropped_rgba)
+
+            # alpha = img[..., 3] > 0
+            H, W = alpha[first_mask].shape
             # get the bounding box of alpha
-            y, x = np.where(alpha)
+            y, x = np.where(alpha[first_mask])
             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
@@ -1102,6 +1122,7 @@ def process_image(input_image, input_text):
             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
@@ -1113,24 +1134,6 @@ def process_image(input_image, input_text):
             image = (image * 255).clip(0, 255).astype(np.uint8)
             image = Image.fromarray(image)
             
-            # 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)
-            # y_min, y_max = y_indices.min(), y_indices.max()
-            # x_min, x_max = x_indices.min(), x_indices.max()
-            
-            # Crop the RGBA image
-            #cropped_rgba = rgba[y_min:y_max+1, x_min:x_max+1]
-            
-            # Set extracted foreground for mask mover
-            # mask_mover.set_extracted_fg(cropped_rgba)
-            
             return annotated_frame, image, gr.update(visible=False), gr.update(visible=False)
 
     
@@ -1244,23 +1247,57 @@ def process_image(input_image, input_text):
             img = input_image.copy()
             H, W, C = img.shape
             
+            first_mask = detections.mask[0]
+            
+
+            
             # Create RGBA image
             alpha = np.zeros((H, W, 1), dtype=np.uint8)
+            
             alpha[first_mask] = 255
-            rgba = np.dstack((img, alpha)).astype(np.uint8)
+            
+            # rgba = np.dstack((img, alpha)).astype(np.uint8)
             
             # Crop to mask bounds to minimize image size
-            y_indices, x_indices = np.where(first_mask)
-            y_min, y_max = y_indices.min(), y_indices.max()
-            x_min, x_max = x_indices.min(), x_indices.max()
+            # y_indices, x_indices = np.where(first_mask)
+            # y_min, y_max = y_indices.min(), y_indices.max()
+            # x_min, x_max = x_indices.min(), x_indices.max()
             
             # Crop the RGBA image
-            cropped_rgba = rgba[y_min:y_max+1, x_min:x_max+1]
+            # cropped_rgba = rgba[y_min:y_max+1, x_min:x_max+1]
             
             # Set extracted foreground for mask mover
             # mask_mover.set_extracted_fg(cropped_rgba)
+
+            # alpha = img[..., 3] > 0
+            H, W = alpha[first_mask].shape
+            # get the bounding box of alpha
+            y, x = np.where(alpha[first_mask])
+            y0, y1 = max(y.min() - 1, 0), min(y.max() + 1, H)
+            x0, x1 = max(x.min() - 1, 0), min(x.max() + 1, W)
             
-            return annotated_frame, cropped_rgba, gr.update(visible=False), gr.update(visible=False)
+            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)
+            
+            return annotated_frame, image, gr.update(visible=False), gr.update(visible=False)
         return annotated_frame, None, gr.update(visible=False), gr.update(visible=False)
         
 
@@ -1317,9 +1354,9 @@ with block:
                         bg_source = gr.Radio(choices=[e.value for e in list(BGSource)[2:]],
                                             value=BGSource.LEFT.value,
                                             label="Lighting Preference (Initial Latent)", type='value')
-                
-                    example_quick_subjects = gr.Dataset(samples=quick_subjects, label='Subject Quick List', samples_per_page=1000, components=[prompt])
-                    example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Lighting Quick List', samples_per_page=1000, components=[prompt])
+                    
+                        example_quick_subjects = gr.Dataset(samples=quick_subjects, label='Subject Quick List', samples_per_page=1000, components=[prompt])
+                        example_quick_prompts = gr.Dataset(samples=quick_prompts, label='Lighting Quick List', samples_per_page=1000, components=[prompt])
                 with gr.Row():    
                     relight_button = gr.Button(value="Relight")