Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| """## hugging face funcs""" | |
| import io | |
| import matplotlib.pyplot as plt | |
| import requests | |
| import inflect | |
| from PIL import Image | |
| def load_image_from_url(url): | |
| return Image.open(requests.get(url, stream=True).raw) | |
| def render_results_in_image(in_pil_img, in_results): | |
| plt.figure(figsize=(16, 10)) | |
| plt.imshow(in_pil_img) | |
| ax = plt.gca() | |
| for prediction in in_results: | |
| x, y = prediction['box']['xmin'], prediction['box']['ymin'] | |
| w = prediction['box']['xmax'] - prediction['box']['xmin'] | |
| h = prediction['box']['ymax'] - prediction['box']['ymin'] | |
| ax.add_patch(plt.Rectangle((x, y), | |
| w, | |
| h, | |
| fill=False, | |
| color="green", | |
| linewidth=2)) | |
| ax.text( | |
| x, | |
| y, | |
| f"{prediction['label']}: {round(prediction['score']*100, 1)}%", | |
| color='red' | |
| ) | |
| plt.axis("off") | |
| # Save the modified image to a BytesIO object | |
| img_buf = io.BytesIO() | |
| plt.savefig(img_buf, format='png', | |
| bbox_inches='tight', | |
| pad_inches=0) | |
| img_buf.seek(0) | |
| modified_image = Image.open(img_buf) | |
| # Close the plot to prevent it from being displayed | |
| plt.close() | |
| return modified_image | |
| def summarize_predictions_natural_language(predictions): | |
| summary = {} | |
| p = inflect.engine() | |
| for prediction in predictions: | |
| label = prediction['label'] | |
| if label in summary: | |
| summary[label] += 1 | |
| else: | |
| summary[label] = 1 | |
| result_string = "In this image, there are " | |
| for i, (label, count) in enumerate(summary.items()): | |
| count_string = p.number_to_words(count) | |
| result_string += f"{count_string} {label}" | |
| if count > 1: | |
| result_string += "s" | |
| result_string += " " | |
| if i == len(summary) - 2: | |
| result_string += "and " | |
| # Remove the trailing comma and space | |
| result_string = result_string.rstrip(', ') + "." | |
| return result_string | |
| ##### To ignore warnings ##### | |
| import warnings | |
| import logging | |
| from transformers import logging as hf_logging | |
| def ignore_warnings(): | |
| # Ignore specific Python warnings | |
| warnings.filterwarnings("ignore", message="Some weights of the model checkpoint") | |
| warnings.filterwarnings("ignore", message="Could not find image processor class") | |
| warnings.filterwarnings("ignore", message="The `max_size` parameter is deprecated") | |
| # Adjust logging for libraries using the logging module | |
| logging.basicConfig(level=logging.ERROR) | |
| hf_logging.set_verbosity_error() | |
| ######## | |
| import numpy as np | |
| import torch | |
| import matplotlib.pyplot as plt | |
| def show_mask(mask, ax, random_color=False): | |
| if random_color: | |
| color = np.concatenate([np.random.random(3), | |
| np.array([0.6])], | |
| axis=0) | |
| else: | |
| color = np.array([30/255, 144/255, 255/255, 0.6]) | |
| h, w = mask.shape[-2:] | |
| mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) | |
| ax.imshow(mask_image) | |
| def show_box(box, ax): | |
| x0, y0 = box[0], box[1] | |
| w, h = box[2] - box[0], box[3] - box[1] | |
| ax.add_patch(plt.Rectangle((x0, y0), | |
| w, | |
| h, edgecolor='green', | |
| facecolor=(0,0,0,0), | |
| lw=2)) | |
| def show_boxes_on_image(raw_image, boxes): | |
| plt.figure(figsize=(10,10)) | |
| plt.imshow(raw_image) | |
| for box in boxes: | |
| show_box(box, plt.gca()) | |
| plt.axis('on') | |
| plt.show() | |
| def show_points_on_image(raw_image, input_points, input_labels=None): | |
| plt.figure(figsize=(10,10)) | |
| plt.imshow(raw_image) | |
| input_points = np.array(input_points) | |
| if input_labels is None: | |
| labels = np.ones_like(input_points[:, 0]) | |
| else: | |
| labels = np.array(input_labels) | |
| show_points(input_points, labels, plt.gca()) | |
| plt.axis('on') | |
| plt.show() | |
| def show_points_and_boxes_on_image(raw_image, | |
| boxes, | |
| input_points, | |
| input_labels=None): | |
| plt.figure(figsize=(10,10)) | |
| plt.imshow(raw_image) | |
| input_points = np.array(input_points) | |
| if input_labels is None: | |
| labels = np.ones_like(input_points[:, 0]) | |
| else: | |
| labels = np.array(input_labels) | |
| show_points(input_points, labels, plt.gca()) | |
| for box in boxes: | |
| show_box(box, plt.gca()) | |
| plt.axis('on') | |
| plt.show() | |
| def show_points_and_boxes_on_image(raw_image, | |
| boxes, | |
| input_points, | |
| input_labels=None): | |
| plt.figure(figsize=(10,10)) | |
| plt.imshow(raw_image) | |
| input_points = np.array(input_points) | |
| if input_labels is None: | |
| labels = np.ones_like(input_points[:, 0]) | |
| else: | |
| labels = np.array(input_labels) | |
| show_points(input_points, labels, plt.gca()) | |
| for box in boxes: | |
| show_box(box, plt.gca()) | |
| plt.axis('on') | |
| plt.show() | |
| def show_points(coords, labels, ax, marker_size=375): | |
| pos_points = coords[labels==1] | |
| neg_points = coords[labels==0] | |
| ax.scatter(pos_points[:, 0], | |
| pos_points[:, 1], | |
| color='green', | |
| marker='*', | |
| s=marker_size, | |
| edgecolor='white', | |
| linewidth=1.25) | |
| ax.scatter(neg_points[:, 0], | |
| neg_points[:, 1], | |
| color='red', | |
| marker='*', | |
| s=marker_size, | |
| edgecolor='white', | |
| linewidth=1.25) | |
| def fig2img(fig): | |
| """Convert a Matplotlib figure to a PIL Image and return it""" | |
| import io | |
| buf = io.BytesIO() | |
| fig.savefig(buf) | |
| buf.seek(0) | |
| img = Image.open(buf) | |
| return img | |
| def show_mask_on_image(raw_image, mask, return_image=False): | |
| if not isinstance(mask, torch.Tensor): | |
| mask = torch.Tensor(mask) | |
| if len(mask.shape) == 4: | |
| mask = mask.squeeze() | |
| fig, axes = plt.subplots(1, 1, figsize=(15, 15)) | |
| mask = mask.cpu().detach() | |
| axes.imshow(np.array(raw_image)) | |
| show_mask(mask, axes) | |
| axes.axis("off") | |
| plt.show() | |
| if return_image: | |
| fig = plt.gcf() | |
| return fig2img(fig) | |
| def show_pipe_masks_on_image(raw_image, outputs, return_image=False): | |
| plt.imshow(np.array(raw_image)) | |
| ax = plt.gca() | |
| for mask in outputs["masks"]: | |
| show_mask(mask, ax=ax, random_color=True) | |
| plt.axis("off") | |
| plt.show() | |
| if return_image: | |
| fig = plt.gcf() | |
| return fig2img(fig) | |
| """## imports""" | |
| from transformers import pipeline | |
| from transformers import SamModel, SamProcessor | |
| from transformers import BlipForImageTextRetrieval | |
| from transformers import AutoProcessor | |
| from transformers.utils import logging | |
| logging.set_verbosity_error() | |
| #ignore_warnings() | |
| import io | |
| import matplotlib.pyplot as plt | |
| import requests | |
| import inflect | |
| from PIL import Image | |
| import os | |
| import gradio as gr | |
| import time | |
| """# Object detection | |
| ## hugging face model ("facebook/detr-resnet-50"). 167MB | |
| """ | |
| od_pipe = pipeline("object-detection", "facebook/detr-resnet-50") | |
| chosen_model = pipeline("object-detection", "hustvl/yolos-small") | |
| """## gradio funcs""" | |
| def get_object_detection_prediction(model_name, raw_image): | |
| model = od_pipe | |
| if "chosen-model" in model_name: | |
| model = chosen_model | |
| start = time.time() | |
| pipeline_output = model(raw_image) | |
| end = time.time() | |
| elapsed_result = f'{model_name} object detection elapsed {end-start} seconds' | |
| print(elapsed_result) | |
| processed_image = render_results_in_image(raw_image, pipeline_output) | |
| return [processed_image, elapsed_result] | |
| """# Image segmentation | |
| ## hugging face models: Zigeng/SlimSAM-uniform-77(segmentation) 39MB, Intel/dpt-hybrid-midas(depth) 490MB | |
| """ | |
| hugging_face_segmentation_pipe = pipeline("mask-generation", "Zigeng/SlimSAM-uniform-77") | |
| hugging_face_segmentation_model = SamModel.from_pretrained("Zigeng/SlimSAM-uniform-77") | |
| hugging_face_segmentation_processor = SamProcessor.from_pretrained("Zigeng/SlimSAM-uniform-77") | |
| hugging_face_depth_estimator = pipeline(task="depth-estimation", model="Intel/dpt-hybrid-midas") | |
| """## chosen models: facebook/sam-vit-base(segmentation) 375MB, LiheYoung/depth-anything-small-hf(depth) 100MB""" | |
| chosen_name = "facebook/sam-vit-base" | |
| chosen_segmentation_pipe = pipeline("mask-generation", chosen_name) | |
| chosen_segmentation_model = SamModel.from_pretrained(chosen_name) | |
| chosen_segmentation_processor = SamProcessor.from_pretrained(chosen_name) | |
| chosen_depth_estimator = pipeline(task="depth-estimation", model="LiheYoung/depth-anything-small-hf") | |
| """## gradio funcs""" | |
| input_points = [[[1600, 700]]] | |
| def segment_image_pretrained(model_name, raw_image): | |
| processor = hugging_face_segmentation_processor | |
| model = hugging_face_segmentation_model | |
| if("chosen" in model_name): | |
| processor = chosen_segmentation_processor | |
| model = chosen_segmentation_model | |
| start = time.time() | |
| inputs = processor(raw_image, | |
| input_points=input_points, | |
| return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| predicted_masks = processor.image_processor.post_process_masks( | |
| outputs.pred_masks, | |
| inputs["original_sizes"], | |
| inputs["reshaped_input_sizes"]) | |
| results = [] | |
| predicted_mask = predicted_masks[0] | |
| end = time.time() | |
| elapsed_result = f'{model_name} pretrained image segmentation elapsed {end-start} seconds' | |
| print(elapsed_result) | |
| for i in range(3): | |
| results.append(show_mask_on_image(raw_image, predicted_mask[:, i], return_image=True)) | |
| results.append(elapsed_result); | |
| return results | |
| def segment_image(model_name, raw_image): | |
| model = hugging_face_segmentation_pipe | |
| if("chosen" in model_name): | |
| print("chosen model used") | |
| model = chosen_segmentation_pipe | |
| start = time.time() | |
| output = model(raw_image, points_per_batch=32) | |
| end = time.time() | |
| elapsed_result = f'{model_name} raw image segmentation elapsed {end-start} seconds' | |
| print(elapsed_result) | |
| return [show_pipe_masks_on_image(raw_image, output, return_image = True), elapsed_result] | |
| def depth_image(model_name, input_image): | |
| depth_estimator = hugging_face_depth_estimator | |
| print(model_name) | |
| if("chosen" in model_name): | |
| print("chosen model used") | |
| depth_estimator = chosen_depth_estimator | |
| start = time.time() | |
| out = depth_estimator(input_image) | |
| prediction = torch.nn.functional.interpolate( | |
| out["predicted_depth"].unsqueeze(0).unsqueeze(0), | |
| size=input_image.size[::-1], | |
| mode="bicubic", | |
| align_corners=False, | |
| ) | |
| end = time.time() | |
| elapsed_result = f'{model_name} Depth Estimation elapsed {end-start} seconds' | |
| print(elapsed_result) | |
| output = prediction.squeeze().numpy() | |
| formatted = (output * 255 / np.max(output)).astype("uint8") | |
| depth = Image.fromarray(formatted) | |
| return [depth, elapsed_result] | |
| """# Image retrieval | |
| ## hugging face model: Salesforce/blip-itm-base-coco 900MB | |
| """ | |
| hugging_face_retrieval_model = BlipForImageTextRetrieval.from_pretrained( | |
| "Salesforce/blip-itm-base-coco") | |
| hugging_face_retrieval_processor = AutoProcessor.from_pretrained( | |
| "Salesforce/blip-itm-base-coco") | |
| """## chosen model: Salesforce/blip-itm-base-flickr 900MB""" | |
| chosen_retrieval_model = BlipForImageTextRetrieval.from_pretrained( | |
| "Salesforce/blip-itm-base-flickr") | |
| chosen_retrieval_processor = AutoProcessor.from_pretrained( | |
| "Salesforce/blip-itm-base-flickr") | |
| """## gradion func""" | |
| def retrieve_image(model_name, raw_image, predict_text): | |
| processor = hugging_face_retrieval_processor | |
| model = hugging_face_retrieval_model | |
| if("chosen" in model_name): | |
| processor = chosen_retrieval_processor | |
| model = chosen_retrieval_model | |
| start = time.time() | |
| inputs = processor(images=raw_image, | |
| text=predict_text, | |
| return_tensors="pt") | |
| end = time.time() | |
| elapsed_result = f"{model_name} image retrieval elapsed {end-start} seconds" | |
| print(elapsed_result) | |
| itm_scores = model(**inputs)[0] | |
| itm_score = torch.nn.functional.softmax(itm_scores,dim=1) | |
| return [f"""\ | |
| The image and text are matched \ | |
| with a probability of {itm_score[0][1]:.4f}""", | |
| elapsed_result] | |
| """# gradio""" | |
| with gr.Blocks() as object_detection_tab: | |
| gr.Markdown("# Detect objects on image") | |
| gr.Markdown("Upload an image, choose model, press button.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| # Input components | |
| input_image = gr.Image(label="Upload Image", type="pil") | |
| model_selector = gr.Dropdown(["hugging-face(facebook/detr-resnet-50)", "chosen-model(hustvl/yolos-small)"], | |
| label = "Select Model") | |
| with gr.Column(): | |
| # Output image | |
| elapsed_result = gr.Textbox(label="Seconds elapsed", lines=1) | |
| output_image = gr.Image(label="Output Image", type="pil") | |
| # Process button | |
| process_btn = gr.Button("Detect objects") | |
| # Connect the input components to the processing function | |
| process_btn.click( | |
| fn=get_object_detection_prediction, | |
| inputs=[ | |
| model_selector, | |
| input_image | |
| ], | |
| outputs=[output_image, elapsed_result] | |
| ) | |
| with gr.Blocks() as image_segmentation_detection_tab: | |
| gr.Markdown("# Image segmentation") | |
| gr.Markdown("Upload an image, choose model, press button.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| # Input components | |
| input_image = gr.Image(label="Upload Image", type="pil") | |
| model_selector = gr.Dropdown(["hugging-face(Zigeng/SlimSAM-uniform-77)", "chosen-model(facebook/sam-vit-base)"], | |
| label = "Select Model") | |
| with gr.Column(): | |
| elapsed_result = gr.Textbox(label="Seconds elapsed", lines=1) | |
| # Output image | |
| output_image = gr.Image(label="Segmented image", type="pil") | |
| with gr.Row(): | |
| with gr.Column(): | |
| segment_btn = gr.Button("Segment image(not pretrained)") | |
| with gr.Row(): | |
| elapsed_result_pretrained_segment = gr.Textbox(label="Seconds elapsed", lines=1) | |
| with gr.Column(): | |
| segment_pretrained_output_image_1 = gr.Image(label="Segmented image by pretrained model", type="pil") | |
| with gr.Column(): | |
| segment_pretrained_output_image_2 = gr.Image(label="Segmented image by pretrained model", type="pil") | |
| with gr.Column(): | |
| segment_pretrained_output_image_3 = gr.Image(label="Segmented image by pretrained model", type="pil") | |
| with gr.Row(): | |
| with gr.Column(): | |
| segment_pretrained_model_selector = gr.Dropdown(["hugging-face(Zigeng/SlimSAM-uniform-77)", "chosen-model(facebook/sam-vit-base)"], | |
| label = "Select Model") | |
| segment_pretrained_btn = gr.Button("Segment image(pretrained)") | |
| with gr.Row(): | |
| with gr.Column(): | |
| depth_output_image = gr.Image(label="Depth image", type="pil") | |
| elapsed_result_depth = gr.Textbox(label="Seconds elapsed", lines=1) | |
| with gr.Row(): | |
| with gr.Column(): | |
| depth_model_selector = gr.Dropdown(["hugging-face(Intel/dpt-hybrid-midas)", "chosen-model(LiheYoung/depth-anything-small-hf)"], | |
| label = "Select Model") | |
| depth_btn = gr.Button("Get image depth") | |
| segment_btn.click( | |
| fn=segment_image, | |
| inputs=[ | |
| model_selector, | |
| input_image | |
| ], | |
| outputs=[output_image, elapsed_result] | |
| ) | |
| segment_pretrained_btn.click( | |
| fn=segment_image_pretrained, | |
| inputs=[ | |
| segment_pretrained_model_selector, | |
| input_image | |
| ], | |
| outputs=[segment_pretrained_output_image_1, segment_pretrained_output_image_2, segment_pretrained_output_image_3, elapsed_result_pretrained_segment] | |
| ) | |
| depth_btn.click( | |
| fn=depth_image, | |
| inputs=[ | |
| depth_model_selector, | |
| input_image, | |
| ], | |
| outputs=[depth_output_image, elapsed_result_depth] | |
| ) | |
| with gr.Blocks() as image_retrieval_tab: | |
| gr.Markdown("# Check is text describes image") | |
| gr.Markdown("Upload an image, choose model, press button.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| # Input components | |
| input_image = gr.Image(label="Upload Image", type="pil") | |
| text_prediction = gr.TextArea(label="Describe image") | |
| model_selector = gr.Dropdown(["hugging-face(Salesforce/blip-itm-base-coco)", "chosen-model(Salesforce/blip-itm-base-flickr)"], | |
| label = "Select Model") | |
| with gr.Column(): | |
| # Output image | |
| output_result = gr.Textbox(label="Probability result", lines=3) | |
| elapsed_result = gr.Textbox(label="Seconds elapsed", lines=1) | |
| # Process button | |
| process_btn = gr.Button("Detect objects") | |
| # Connect the input components to the processing function | |
| process_btn.click( | |
| fn=retrieve_image, | |
| inputs=[ | |
| model_selector, | |
| input_image, | |
| text_prediction | |
| ], | |
| outputs=[output_result, elapsed_result] | |
| ) | |
| with gr.Blocks() as app: | |
| gr.TabbedInterface( | |
| [object_detection_tab, | |
| image_segmentation_detection_tab, | |
| image_retrieval_tab], | |
| ["Object detection", | |
| "Image segmentation", | |
| "Retrieve image" | |
| ], | |
| ) | |
| app.launch(share=True, debug=True) | |
| app.close() |