diff --git a/.gitignore b/.gitignore index ce87ac5..bddefc7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ /venv -/.idea \ No newline at end of file +/.idea +/tmp \ No newline at end of file diff --git a/app.py b/app.py index 3f43163..ec9f684 100644 --- a/app.py +++ b/app.py @@ -1,16 +1,22 @@ +import os from typing import Tuple, Optional +import cv2 import gradio as gr +import numpy as np +import spaces import supervision as sv import torch from PIL import Image +from tqdm import tqdm +from utils.video import generate_unique_name, create_directory, delete_directory from utils.florence import load_florence_model, run_florence_inference, \ FLORENCE_DETAILED_CAPTION_TASK, \ FLORENCE_CAPTION_TO_PHRASE_GROUNDING_TASK, FLORENCE_OPEN_VOCABULARY_DETECTION_TASK -from utils.modes import INFERENCE_MODES, OPEN_VOCABULARY_DETECTION, \ - CAPTION_GROUNDING_MASKS -from utils.sam import load_sam_model, run_sam_inference +from utils.modes import IMAGE_INFERENCE_MODES, IMAGE_OPEN_VOCABULARY_DETECTION_MODE, \ + IMAGE_CAPTION_GROUNDING_MASKS_MODE, VIDEO_INFERENCE_MODES +from utils.sam import load_sam_image_model, run_sam_inference, load_sam_video_model MARKDOWN = """ # Florence2 + SAM2 🔥 @@ -33,29 +39,46 @@ MARKDOWN = """ This demo integrates Florence2 and SAM2 by creating a two-stage inference pipeline. In the first stage, Florence2 performs tasks such as object detection, open-vocabulary object detection, image captioning, or phrase grounding. In the second stage, SAM2 -performs object segmentation on the image. **Video segmentation will be available -soon.** +performs object segmentation on the image. """ -EXAMPLES = [ - [OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", 'straw'], - [OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", 'napkin'], - [OPEN_VOCABULARY_DETECTION, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", 'tail'], - [CAPTION_GROUNDING_MASKS, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None], - [CAPTION_GROUNDING_MASKS, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", None], +IMAGE_PROCESSING_EXAMPLES = [ + [IMAGE_OPEN_VOCABULARY_DETECTION_MODE, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", 'straw, white napkin, black napkin, dog, hair, man'], + [IMAGE_OPEN_VOCABULARY_DETECTION_MODE, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", 'tail'], + [IMAGE_CAPTION_GROUNDING_MASKS_MODE, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None], + [IMAGE_CAPTION_GROUNDING_MASKS_MODE, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", None], ] +VIDEO_SCALE_FACTOR = 0.5 +VIDEO_TARGET_DIRECTORY = "tmp" +create_directory(directory_path=VIDEO_TARGET_DIRECTORY) + DEVICE = torch.device("cuda") +# DEVICE = torch.device("cpu") + +torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__() +if torch.cuda.get_device_properties(0).major >= 8: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + FLORENCE_MODEL, FLORENCE_PROCESSOR = load_florence_model(device=DEVICE) -SAM_MODEL = load_sam_model(device=DEVICE) -BOX_ANNOTATOR = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX) +SAM_IMAGE_MODEL = load_sam_image_model(device=DEVICE) +SAM_VIDEO_MODEL = load_sam_video_model(device=DEVICE) +COLORS = ['#FF1493', '#00BFFF', '#FF6347', '#FFD700', '#32CD32', '#8A2BE2'] +COLOR_PALETTE = sv.ColorPalette.from_hex(COLORS) +BOX_ANNOTATOR = sv.BoxAnnotator(color=COLOR_PALETTE, color_lookup=sv.ColorLookup.INDEX) LABEL_ANNOTATOR = sv.LabelAnnotator( + color=COLOR_PALETTE, color_lookup=sv.ColorLookup.INDEX, text_position=sv.Position.CENTER_OF_MASS, - text_color=sv.Color.from_hex("#FFFFFF"), + text_color=sv.Color.from_hex("#000000"), border_radius=5 ) -MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX) +MASK_ANNOTATOR = sv.MaskAnnotator( + color=COLOR_PALETTE, + color_lookup=sv.ColorLookup.INDEX +) def annotate_image(image, detections): @@ -68,38 +91,50 @@ def annotate_image(image, detections): def on_mode_dropdown_change(text): return [ - gr.Textbox(visible=text == OPEN_VOCABULARY_DETECTION), - gr.Textbox(visible=text == CAPTION_GROUNDING_MASKS), + gr.Textbox(visible=text == IMAGE_OPEN_VOCABULARY_DETECTION_MODE), + gr.Textbox(visible=text == IMAGE_CAPTION_GROUNDING_MASKS_MODE), ] -def process( +@spaces.GPU +@torch.inference_mode() +@torch.autocast(device_type="cuda", dtype=torch.bfloat16) +def process_image( mode_dropdown, image_input, text_input ) -> Tuple[Optional[Image.Image], Optional[str]]: if not image_input: + gr.Info("Please upload an image.") return None, None - if mode_dropdown == OPEN_VOCABULARY_DETECTION: + if mode_dropdown == IMAGE_OPEN_VOCABULARY_DETECTION_MODE: if not text_input: + gr.Info("Please enter a text prompt.") return None, None - _, result = run_florence_inference( - model=FLORENCE_MODEL, - processor=FLORENCE_PROCESSOR, - device=DEVICE, - image=image_input, - task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK, - text=text_input - ) - detections = sv.Detections.from_lmm( - lmm=sv.LMM.FLORENCE_2, - result=result, - resolution_wh=image_input.size - ) - detections = run_sam_inference(SAM_MODEL, image_input, detections) + texts = [prompt.strip() for prompt in text_input.split(",")] + detections_list = [] + for text in texts: + _, result = run_florence_inference( + model=FLORENCE_MODEL, + processor=FLORENCE_PROCESSOR, + device=DEVICE, + image=image_input, + task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK, + text=text + ) + detections = sv.Detections.from_lmm( + lmm=sv.LMM.FLORENCE_2, + result=result, + resolution_wh=image_input.size + ) + detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections) + detections_list.append(detections) + + detections = sv.Detections.merge(detections_list) + detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections) return annotate_image(image_input, detections), None - if mode_dropdown == CAPTION_GROUNDING_MASKS: + if mode_dropdown == IMAGE_CAPTION_GROUNDING_MASKS_MODE: _, result = run_florence_inference( model=FLORENCE_MODEL, processor=FLORENCE_PROCESSOR, @@ -121,65 +156,228 @@ def process( result=result, resolution_wh=image_input.size ) - detections = run_sam_inference(SAM_MODEL, image_input, detections) + detections = run_sam_inference(SAM_IMAGE_MODEL, image_input, detections) return annotate_image(image_input, detections), caption +@spaces.GPU(duration=300) +@torch.inference_mode() +@torch.autocast(device_type="cuda", dtype=torch.bfloat16) +def process_video( + mode_dropdown, video_input, text_input, progress=gr.Progress(track_tqdm=True) +) -> Optional[str]: + if not video_input: + gr.Info("Please upload a video.") + return None + + if not text_input: + gr.Info("Please enter a text prompt.") + return None + + frame_generator = sv.get_video_frames_generator(video_input) + frame = next(frame_generator) + frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + + texts = [prompt.strip() for prompt in text_input.split(",")] + detections_list = [] + for text in texts: + _, result = run_florence_inference( + model=FLORENCE_MODEL, + processor=FLORENCE_PROCESSOR, + device=DEVICE, + image=frame, + task=FLORENCE_OPEN_VOCABULARY_DETECTION_TASK, + text=text + ) + detections = sv.Detections.from_lmm( + lmm=sv.LMM.FLORENCE_2, + result=result, + resolution_wh=frame.size + ) + detections = run_sam_inference(SAM_IMAGE_MODEL, frame, detections) + detections_list.append(detections) + + detections = sv.Detections.merge(detections_list) + detections = run_sam_inference(SAM_IMAGE_MODEL, frame, detections) + + if len(detections.mask) == 0: + gr.Info( + "No objects of class {text_input} found in the first frame of the video. " + "Trim the video to make the object appear in the first frame or try a " + "different text prompt." + ) + return None + + name = generate_unique_name() + frame_directory_path = os.path.join(VIDEO_TARGET_DIRECTORY, name) + frames_sink = sv.ImageSink( + target_dir_path=frame_directory_path, + image_name_pattern="{:05d}.jpeg" + ) + + video_info = sv.VideoInfo.from_video_path(video_input) + video_info.width = int(video_info.width * VIDEO_SCALE_FACTOR) + video_info.height = int(video_info.height * VIDEO_SCALE_FACTOR) + + frames_generator = sv.get_video_frames_generator(video_input) + with frames_sink: + for frame in tqdm( + frames_generator, + total=video_info.total_frames, + desc="splitting video into frames" + ): + frame = sv.scale_image(frame, VIDEO_SCALE_FACTOR) + frames_sink.save_image(frame) + + inference_state = SAM_VIDEO_MODEL.init_state( + video_path=frame_directory_path, + device=DEVICE + ) + + for mask_index, mask in enumerate(detections.mask): + _, object_ids, mask_logits = SAM_VIDEO_MODEL.add_new_mask( + inference_state=inference_state, + frame_idx=0, + obj_id=mask_index, + mask=mask + ) + + video_path = os.path.join(VIDEO_TARGET_DIRECTORY, f"{name}.mp4") + frames_generator = sv.get_video_frames_generator(video_input) + masks_generator = SAM_VIDEO_MODEL.propagate_in_video(inference_state) + with sv.VideoSink(video_path, video_info=video_info) as sink: + for frame, (_, tracker_ids, mask_logits) in zip(frames_generator, masks_generator): + frame = sv.scale_image(frame, VIDEO_SCALE_FACTOR) + masks = (mask_logits > 0.0).cpu().numpy().astype(bool) + if len(masks.shape) == 4: + masks = np.squeeze(masks, axis=1) + + detections = sv.Detections( + xyxy=sv.mask_to_xyxy(masks=masks), + mask=masks, + class_id=np.array(tracker_ids) + ) + annotated_frame = frame.copy() + annotated_frame = MASK_ANNOTATOR.annotate( + scene=annotated_frame, detections=detections) + annotated_frame = BOX_ANNOTATOR.annotate( + scene=annotated_frame, detections=detections) + sink.write_frame(annotated_frame) + + delete_directory(frame_directory_path) + return video_path + + with gr.Blocks() as demo: gr.Markdown(MARKDOWN) - mode_dropdown_component = gr.Dropdown( - choices=INFERENCE_MODES, - value=INFERENCE_MODES[0], - label="Mode", - info="Select a mode to use.", - interactive=True - ) - with gr.Row(): - with gr.Column(): - image_input_component = gr.Image( - type='pil', label='Upload image') - text_input_component = gr.Textbox( - label='Text prompt') - submit_button_component = gr.Button(value='Submit', variant='primary') - with gr.Column(): - image_output_component = gr.Image(type='pil', label='Image output') - text_output_component = gr.Textbox(label='Caption output', visible=False) - - with gr.Row(): - gr.Examples( - fn=process, - examples=EXAMPLES, - inputs=[ - mode_dropdown_component, - image_input_component, - text_input_component - ], - outputs=[ - image_output_component, - text_output_component - ], - run_on_click=True + with gr.Tab("Image"): + image_processing_mode_dropdown_component = gr.Dropdown( + choices=IMAGE_INFERENCE_MODES, + value=IMAGE_INFERENCE_MODES[0], + label="Mode", + info="Select a mode to use.", + interactive=True ) + with gr.Row(): + with gr.Column(): + image_processing_image_input_component = gr.Image( + type='pil', label='Upload image') + image_processing_text_input_component = gr.Textbox( + label='Text prompt', + placeholder='Enter comma separated text prompts') + image_processing_submit_button_component = gr.Button( + value='Submit', variant='primary') + with gr.Column(): + image_processing_image_output_component = gr.Image( + type='pil', label='Image output') + image_processing_text_output_component = gr.Textbox( + label='Caption output', visible=False) - submit_button_component.click( - fn=process, + with gr.Row(): + gr.Examples( + fn=process_image, + examples=IMAGE_PROCESSING_EXAMPLES, + inputs=[ + image_processing_mode_dropdown_component, + image_processing_image_input_component, + image_processing_text_input_component + ], + outputs=[ + image_processing_image_output_component, + image_processing_text_output_component + ], + run_on_click=True + ) + with gr.Tab("Video"): + video_processing_mode_dropdown_component = gr.Dropdown( + choices=VIDEO_INFERENCE_MODES, + value=VIDEO_INFERENCE_MODES[0], + label="Mode", + info="Select a mode to use.", + interactive=True + ) + with gr.Row(): + with gr.Column(): + video_processing_video_input_component = gr.Video( + label='Upload video') + video_processing_text_input_component = gr.Textbox( + label='Text prompt', + placeholder='Enter comma separated text prompts') + video_processing_submit_button_component = gr.Button( + value='Submit', variant='primary') + with gr.Column(): + video_processing_video_output_component = gr.Video( + label='Video output') + + image_processing_submit_button_component.click( + fn=process_image, inputs=[ - mode_dropdown_component, - image_input_component, - text_input_component + image_processing_mode_dropdown_component, + image_processing_image_input_component, + image_processing_text_input_component ], outputs=[ - image_output_component, - text_output_component + image_processing_image_output_component, + image_processing_text_output_component ] ) - mode_dropdown_component.change( - on_mode_dropdown_change, - inputs=[mode_dropdown_component], + image_processing_text_input_component.submit( + fn=process_image, + inputs=[ + image_processing_mode_dropdown_component, + image_processing_image_input_component, + image_processing_text_input_component + ], outputs=[ - text_input_component, - text_output_component + image_processing_image_output_component, + image_processing_text_output_component ] ) + image_processing_mode_dropdown_component.change( + on_mode_dropdown_change, + inputs=[image_processing_mode_dropdown_component], + outputs=[ + image_processing_text_input_component, + image_processing_text_output_component + ] + ) + video_processing_submit_button_component.click( + fn=process_video, + inputs=[ + video_processing_mode_dropdown_component, + video_processing_video_input_component, + video_processing_text_input_component + ], + outputs=video_processing_video_output_component + ) + video_processing_text_input_component.submit( + fn=process_video, + inputs=[ + video_processing_mode_dropdown_component, + video_processing_video_input_component, + video_processing_text_input_component + ], + outputs=video_processing_video_output_component + ) demo.launch(debug=False, show_error=True) diff --git a/requirements-local.txt b/requirements-local.txt deleted file mode 100644 index 2371f73..0000000 --- a/requirements-local.txt +++ /dev/null @@ -1,10 +0,0 @@ -torch -einops -spaces -timm -transformers -samv2 -gradio -supervision -opencv-python -pytest \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index de402a5..e2d95e2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +tqdm einops spaces timm diff --git a/utils/modes.py b/utils/modes.py index fb1db4b..008e098 100644 --- a/utils/modes.py +++ b/utils/modes.py @@ -1,7 +1,13 @@ -OPEN_VOCABULARY_DETECTION = "open vocabulary detection + masks" -CAPTION_GROUNDING_MASKS = "caption + grounding + masks" +IMAGE_OPEN_VOCABULARY_DETECTION_MODE = "open vocabulary detection + image masks" +IMAGE_CAPTION_GROUNDING_MASKS_MODE = "caption + grounding + image masks" -INFERENCE_MODES = [ - OPEN_VOCABULARY_DETECTION, - CAPTION_GROUNDING_MASKS +IMAGE_INFERENCE_MODES = [ + IMAGE_OPEN_VOCABULARY_DETECTION_MODE, + IMAGE_CAPTION_GROUNDING_MASKS_MODE +] + +VIDEO_OPEN_VOCABULARY_DETECTION_MODE = "open vocabulary detection + video masks" + +VIDEO_INFERENCE_MODES = [ + VIDEO_OPEN_VOCABULARY_DETECTION_MODE ] diff --git a/utils/sam.py b/utils/sam.py index 74ebad2..16cde29 100644 --- a/utils/sam.py +++ b/utils/sam.py @@ -4,14 +4,14 @@ import numpy as np import supervision as sv import torch from PIL import Image -from sam2.build_sam import build_sam2 +from sam2.build_sam import build_sam2, build_sam2_video_predictor from sam2.sam2_image_predictor import SAM2ImagePredictor SAM_CHECKPOINT = "checkpoints/sam2_hiera_small.pt" SAM_CONFIG = "sam2_hiera_s.yaml" -def load_sam_model( +def load_sam_image_model( device: torch.device, config: str = SAM_CONFIG, checkpoint: str = SAM_CHECKPOINT @@ -20,6 +20,14 @@ def load_sam_model( return SAM2ImagePredictor(sam_model=model) +def load_sam_video_model( + device: torch.device, + config: str = SAM_CONFIG, + checkpoint: str = SAM_CHECKPOINT +) -> Any: + return build_sam2_video_predictor(config, checkpoint, device=device) + + def run_sam_inference( model: Any, image: Image, diff --git a/utils/video.py b/utils/video.py new file mode 100644 index 0000000..ec99fc1 --- /dev/null +++ b/utils/video.py @@ -0,0 +1,26 @@ +import datetime +import os +import shutil +import uuid + + +def create_directory(directory_path: str) -> None: + if not os.path.exists(directory_path): + os.makedirs(directory_path) + + +def delete_directory(directory_path: str) -> None: + if not os.path.exists(directory_path): + raise FileNotFoundError(f"Directory '{directory_path}' does not exist.") + + try: + shutil.rmtree(directory_path) + except PermissionError: + raise PermissionError( + f"Permission denied: Unable to delete '{directory_path}'.") + + +def generate_unique_name(): + current_datetime = datetime.datetime.now().strftime("%Y%m%d%H%M%S") + unique_id = uuid.uuid4() + return f"{current_datetime}_{unique_id}"