Remove commented code

app.py

@@ -32,7 +32,7 @@ import torchvision
 import gradio as gr
 import argparse
 from PIL import Image, ImageDraw, ImageFont
-from scipy import ndimage
+# from scipy import ndimage
 
 # Grounding DINO
 import GroundingDINO.groundingdino.datasets.transforms as T
@@ -104,17 +104,7 @@ class ModelManager:
                 _models['sam_predictor'] = SamPredictor(sam)
                 
                 logger.info(f"SAM-HQ model loaded in {time.time() - start_time:.2f} seconds")
-                
-            # elif model_name == 'blip' and (_models['blip_processor'] is None or _models['blip_model'] is None):
-            #     logger.info("Loading BLIP model...")
-            #     start_time = time.time()
-                
-            #     _models['blip_processor'] = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
-            #     _models['blip_model'] = BlipForConditionalGeneration.from_pretrained(
-            #         "Salesforce/blip-image-captioning-large", torch_dtype=torch.float16
-            #     ).to(device)
-                
-            #     logger.info(f"BLIP model loaded in {time.time() - start_time:.2f} seconds")
+
                 
         except Exception as e:
             logger.error(f"Error loading {model_name} model: {e}")
@@ -137,22 +127,6 @@ class ModelManager:
                 torch.cuda.empty_cache()
 
 
-# def generate_caption(raw_image: Image.Image) -> str:
-#     """Generate image caption using BLIP"""
-#     try:
-#         blip_processor = ModelManager.get_model('blip_processor')
-#         blip_model = ModelManager.get_model('blip_model')
-        
-#         inputs = blip_processor(raw_image, return_tensors="pt").to(device, torch.float16)
-#         out = blip_model.generate(**inputs)
-#         caption = blip_processor.decode(out[0], skip_special_tokens=True)
-#         logger.info(f"Generated caption: {caption}")
-#         return caption
-#     except Exception as e:
-#         logger.error(f"Error generating caption: {e}")
-#         return "Failed to generate caption."
-
-
 def transform_image(image_pil: Image.Image) -> torch.Tensor:
     """Transform PIL image for GroundingDINO"""
     transform = T.Compose([
@@ -218,10 +192,7 @@ def get_grounding_output(
 
 def draw_mask(mask: np.ndarray, draw: ImageDraw.Draw) -> None:
     """Draw mask on image"""
-    # if random_color:
-    #     color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), 153)
-    # else:
-    #     color = (30, 144, 255, 153)
+
     color = (255, 255, 255, 255)
 
     nonzero_coords = np.transpose(np.nonzero(mask))
@@ -245,46 +216,6 @@ def draw_box(box: torch.Tensor, draw: ImageDraw.Draw, label: Optional[str]) -> N
         draw.text((box[0], box[1]), str(label), fill="white")
 
 
-# def draw_point(point: np.ndarray, draw: ImageDraw.Draw, r: int = 10) -> None:
-#     """Draw points on image"""
-#     for p in point:
-#         x, y = p
-#         draw.ellipse((x-r, y-r, x+r, y+r), fill='green')
-
-
-# def process_scribble_points(scribble: np.ndarray) -> np.ndarray:
-#     """Process scribble mask to get point coordinates"""
-#     # Transpose to get the correct orientation
-#     scribble = scribble.transpose(2, 1, 0)[0]
-#     # Label connected components
-#     labeled_array, num_features = ndimage.label(scribble >= 255)
-#     if num_features == 0:
-#         logger.warning("No points detected in scribble")
-#         return np.array([])
-    
-#     # Get center of mass for each component
-#     centers = ndimage.center_of_mass(scribble, labeled_array, range(1, num_features + 1))
-#     return np.array(centers)
-
-
-# def process_scribble_box(scribble: np.ndarray) -> torch.Tensor:
-#     """Process scribble mask to get bounding box"""
-#     # Get point coordinates first
-#     centers = process_scribble_points(scribble)
-#     if len(centers) < 2:
-#         logger.warning("Not enough points for bounding box, need at least 2")
-#         # Return a default small box in the center if not enough points
-#         return torch.tensor([[0.4, 0.4, 0.6, 0.6]])
-    
-#     # Define bounding box from scribble centers: (x_min, y_min, x_max, y_max)
-#     x_min = centers[:, 0].min()
-#     x_max = centers[:, 0].max()
-#     y_min = centers[:, 1].min()
-#     y_max = centers[:, 1].max()
-#     bbox = np.array([x_min, y_min, x_max, y_max])
-#     return torch.tensor(bbox).unsqueeze(0)
-
-
 def run_grounded_sam(
     input_image
     # text_prompt: str, 
@@ -318,13 +249,7 @@ def run_grounded_sam(
         if image_pil is None:
             logger.error("No input image provided")
             return [Image.new('RGB', (400, 300), color='gray')]
-        
-        # # Prepare for scribble tasks
-        # if task_type == 'scribble_box' or task_type == 'scribble_point':
-        #     if scribble is None:
-        #         logger.warning(f"No scribble provided for {task_type} task")
-        #         scribble = np.zeros((image_pil.height, image_pil.width, 3), dtype=np.uint8)
-        
+
         # Transform image for GroundingDINO
         transformed_image = transform_image(image_pil)
         
@@ -337,49 +262,7 @@ def run_grounded_sam(
         boxes_filt, scores, pred_phrases = get_grounding_output(
             transformed_image, text_prompt, box_threshold, text_threshold
         )
-                
-        # # Process based on task type
-        # if task_type == 'automatic':
-        #     # Generate caption with BLIP
-        #     ModelManager.load_model('blip')
-        #     text_prompt = generate_caption(image_pil)
-        #     logger.info(f"Automatic caption: {text_prompt}")
-            
-        #     # Run GroundingDINO
-        #     boxes_filt, scores, pred_phrases = get_grounding_output(
-        #         transformed_image, text_prompt, box_threshold, text_threshold
-        #     )
-            
-        # elif task_type == 'text':
-        #     if not text_prompt:
-        #         logger.warning("No text prompt provided for 'text' task")
-        #         return [image_pil, Image.new('RGBA', size, color=(0, 0, 0, 0))]
-            
-        #     # Run GroundingDINO with provided text
-        #     boxes_filt, scores, pred_phrases = get_grounding_output(
-        #         transformed_image, text_prompt, box_threshold, text_threshold
-        #     )
-            
-        # elif task_type == 'scribble_box':
-        #     # No need for GroundingDINO, get box from scribble
-        #     boxes_filt = process_scribble_box(scribble)
-        #     scores = torch.ones(boxes_filt.size(0))
-        #     pred_phrases = ["scribble_box"] * boxes_filt.size(0)
-            
-        # elif task_type == 'scribble_point':
-        #     # Will handle differently with SAM
-        #     point_coords = process_scribble_points(scribble)
-        #     if len(point_coords) == 0:
-        #         logger.warning("No points detected in scribble")
-        #         return [image_pil, Image.new('RGBA', size, color=(0, 0, 0, 0))]
-            
-        #     boxes_filt = None  # Not needed for point-based segmentation
-            
-        # else:
-        #     logger.error(f"Unknown task type: {task_type}")
-        #     return [image_pil, Image.new('RGBA', size, color=(0, 0, 0, 0))]
-        
-        # Process boxes if present (not for scribble_point)
+
         if boxes_filt is not None:
             # Scale boxes to image dimensions
             for i in range(boxes_filt.size(0)):
@@ -403,10 +286,6 @@ def run_grounded_sam(
         image = np.array(image_pil)
         sam_predictor.set_image(image)
         
-        # # Convert string to boolean
-        # if isinstance(hq_token_only, str):
-        #     hq_token_only = (hq_token_only.lower() == 'true')
-        
         # Run SAM 
         # Use boxes for these task types
         if boxes_filt.size(0) == 0:
@@ -422,55 +301,21 @@ def run_grounded_sam(
             multimask_output=False,
             hq_token_only=hq_token_only,
         )
-            
-        # elif task_type == 'scribble_point':
-        #     # Use points for this task type
-        #     point_labels = np.ones(point_coords.shape[0])
-            
-        #     masks, _, _ = sam_predictor.predict(
-        #         point_coords=point_coords,
-        #         point_labels=point_labels,
-        #         box=None,
-        #         multimask_output=False,
-        #         hq_token_only=hq_token_only,
-        #     )
         
         # Create mask image
         mask_image = Image.new('RGBA', size, color=(0, 0, 0, 0))
         mask_draw = ImageDraw.Draw(mask_image)
         
         # Draw masks
-        if task_type == 'text':
-        #     for mask in masks:
-        #         draw_mask(mask, mask_draw, random_color=True)
-        # else:
-            for mask in masks:
-                draw_mask(mask[0].cpu().numpy(), mask_draw)
+        for mask in masks:
+            draw_mask(mask[0].cpu().numpy(), mask_draw)
         
         # Draw boxes and points on original image
         image_draw = ImageDraw.Draw(image_pil)
 
         for box, label in zip(boxes_filt, pred_phrases):
             draw_box(box, image_draw, label)
-        
-        # if task_type == 'scribble_box':
-        #     for box in boxes_filt:
-        #         draw_box(box, image_draw, None)
-        # elif task_type in ['text', 'automatic']:
-        #     for box, label in zip(boxes_filt, pred_phrases):
-        #         draw_box(box, image_draw, label)
-        # elif task_type == 'scribble_point':
-        #     draw_point(point_coords, image_draw)
-        
-        # Add caption text for automatic mode
-        # if task_type == 'automatic':
-        #     image_draw.text((10, 10), text_prompt, fill='black')
-        
-        # Combine original image with mask
-        # image_pil = image_pil.convert('RGBA')
-        # image_pil.alpha_composite(mask_image)
-        
-        # return [image_pil, mask_image]
+
         return [mask_image]
     
     except Exception as e:
@@ -494,64 +339,23 @@ def create_ui():
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(type="pil", label="image")
-                # input_image = gr.ImageMask(
-                #     sources=["upload", "clipboard"],
-                #     transforms=[],
-                #     layers=False,
-                #     format="pil",
-                #     label="base image",
-                #     show_label=True
-                # )
-                # task_type = gr.Dropdown(
-                #     ["automatic", "scribble_point", "scribble_box", "text"], 
-                #     value="automatic", 
-                #     label="Task Type"
-                # )
-                # text_prompt = gr.Textbox(label="Text Prompt", placeholder="bench .")
-                # hq_token_only = gr.Dropdown(
-                #     [False, True], value=False, label="hq_token_only"
-                # )
+
                 run_button = gr.Button(value='Run')
-                
-                # with gr.Accordion("Advanced options", open=False):
-                #     box_threshold = gr.Slider(
-                #         label="Box Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.001
-                #     )
-                #     text_threshold = gr.Slider(
-                #         label="Text Threshold", minimum=0.0, maximum=1.0, value=0.25, step=0.001
-                #     )
-                #     iou_threshold = gr.Slider(
-                #         label="IOU Threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.001
-                #     )
 
             with gr.Column():
                 gallery = gr.Gallery(
                     label="Generated images", show_label=False, elem_id="gallery"
                 )
-
-        # # Update visibility of text prompt based on task type
-        # def update_text_prompt_visibility(task):
-        #     return gr.update(visible=(task == "text"))
-        
-        # task_type.change(
-        #     fn=update_text_prompt_visibility,
-        #     inputs=[task_type],
-        #     outputs=[text_prompt]
-        # )
         
         # Run button
         run_button.click(
             fn=run_grounded_sam, 
             inputs=[
                 input_image
-                # , text_prompt, task_type, 
-                # box_threshold, text_threshold, iou_threshold, hq_token_only
             ], 
             outputs=gallery
-        )
-        
-        
-    
+        )     
+ 
     return block