prepconcede/image-processing
0
1# -*- coding: utf-8 -*-2 3 4"""## hugging face funcs"""5 6import io7import matplotlib.pyplot as plt8import requests9import inflect10from PIL import Image11 12def load_image_from_url(url):13 return Image.open(requests.get(url, stream=True).raw)14 15def render_results_in_image(in_pil_img, in_results):16 plt.figure(figsize=(16, 10))17 plt.imshow(in_pil_img)18 19 ax = plt.gca()20 21 for prediction in in_results:22 23 x, y = prediction['box']['xmin'], prediction['box']['ymin']24 w = prediction['box']['xmax'] - prediction['box']['xmin']25 h = prediction['box']['ymax'] - prediction['box']['ymin']26 27 ax.add_patch(plt.Rectangle((x, y),28 w,29 h,30 fill=False,31 color="green",32 linewidth=2))33 ax.text(34 x,35 y,36 f"{prediction['label']}: {round(prediction['score']*100, 1)}%",37 color='red'38 )39 40 plt.axis("off")41 42 # Save the modified image to a BytesIO object43 img_buf = io.BytesIO()44 plt.savefig(img_buf, format='png',45 bbox_inches='tight',46 pad_inches=0)47 img_buf.seek(0)48 modified_image = Image.open(img_buf)49 50 # Close the plot to prevent it from being displayed51 plt.close()52 53 return modified_image54 55def summarize_predictions_natural_language(predictions):56 summary = {}57 p = inflect.engine()58 59 for prediction in predictions:60 label = prediction['label']61 if label in summary:62 summary[label] += 163 else:64 summary[label] = 165 66 result_string = "In this image, there are "67 for i, (label, count) in enumerate(summary.items()):68 count_string = p.number_to_words(count)69 result_string += f"{count_string} {label}"70 if count > 1:71 result_string += "s"72 73 result_string += " "74 75 if i == len(summary) - 2:76 result_string += "and "77 78 # Remove the trailing comma and space79 result_string = result_string.rstrip(', ') + "."80 81 return result_string82 83 84##### To ignore warnings #####85import warnings86import logging87from transformers import logging as hf_logging88 89def ignore_warnings():90 # Ignore specific Python warnings91 warnings.filterwarnings("ignore", message="Some weights of the model checkpoint")92 warnings.filterwarnings("ignore", message="Could not find image processor class")93 warnings.filterwarnings("ignore", message="The `max_size` parameter is deprecated")94 95 # Adjust logging for libraries using the logging module96 logging.basicConfig(level=logging.ERROR)97 hf_logging.set_verbosity_error()98 99########100 101import numpy as np102import torch103import matplotlib.pyplot as plt104 105 106def show_mask(mask, ax, random_color=False):107 if random_color:108 color = np.concatenate([np.random.random(3),109 np.array([0.6])],110 axis=0)111 else:112 color = np.array([30/255, 144/255, 255/255, 0.6])113 h, w = mask.shape[-2:]114 mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)115 ax.imshow(mask_image)116 117 118def show_box(box, ax):119 x0, y0 = box[0], box[1]120 w, h = box[2] - box[0], box[3] - box[1]121 ax.add_patch(plt.Rectangle((x0, y0),122 w,123 h, edgecolor='green',124 facecolor=(0,0,0,0),125 lw=2))126 127def show_boxes_on_image(raw_image, boxes):128 plt.figure(figsize=(10,10))129 plt.imshow(raw_image)130 for box in boxes:131 show_box(box, plt.gca())132 plt.axis('on')133 plt.show()134 135def show_points_on_image(raw_image, input_points, input_labels=None):136 plt.figure(figsize=(10,10))137 plt.imshow(raw_image)138 input_points = np.array(input_points)139 if input_labels is None:140 labels = np.ones_like(input_points[:, 0])141 else:142 labels = np.array(input_labels)143 show_points(input_points, labels, plt.gca())144 plt.axis('on')145 plt.show()146 147def show_points_and_boxes_on_image(raw_image,148 boxes,149 input_points,150 input_labels=None):151 plt.figure(figsize=(10,10))152 plt.imshow(raw_image)153 input_points = np.array(input_points)154 if input_labels is None:155 labels = np.ones_like(input_points[:, 0])156 else:157 labels = np.array(input_labels)158 show_points(input_points, labels, plt.gca())159 for box in boxes:160 show_box(box, plt.gca())161 plt.axis('on')162 plt.show()163 164 165def show_points_and_boxes_on_image(raw_image,166 boxes,167 input_points,168 input_labels=None):169 plt.figure(figsize=(10,10))170 plt.imshow(raw_image)171 input_points = np.array(input_points)172 if input_labels is None:173 labels = np.ones_like(input_points[:, 0])174 else:175 labels = np.array(input_labels)176 show_points(input_points, labels, plt.gca())177 for box in boxes:178 show_box(box, plt.gca())179 plt.axis('on')180 plt.show()181 182 183def show_points(coords, labels, ax, marker_size=375):184 pos_points = coords[labels==1]185 neg_points = coords[labels==0]186 ax.scatter(pos_points[:, 0],187 pos_points[:, 1],188 color='green',189 marker='*',190 s=marker_size,191 edgecolor='white',192 linewidth=1.25)193 ax.scatter(neg_points[:, 0],194 neg_points[:, 1],195 color='red',196 marker='*',197 s=marker_size,198 edgecolor='white',199 linewidth=1.25)200 201 202def fig2img(fig):203 """Convert a Matplotlib figure to a PIL Image and return it"""204 import io205 buf = io.BytesIO()206 fig.savefig(buf)207 buf.seek(0)208 img = Image.open(buf)209 return img210 211 212def show_mask_on_image(raw_image, mask, return_image=False):213 if not isinstance(mask, torch.Tensor):214 mask = torch.Tensor(mask)215 216 if len(mask.shape) == 4:217 mask = mask.squeeze()218 219 fig, axes = plt.subplots(1, 1, figsize=(15, 15))220 221 mask = mask.cpu().detach()222 axes.imshow(np.array(raw_image))223 show_mask(mask, axes)224 axes.axis("off")225 plt.show()226 227 if return_image:228 fig = plt.gcf()229 return fig2img(fig)230 231 232 233 234def show_pipe_masks_on_image(raw_image, outputs, return_image=False):235 plt.imshow(np.array(raw_image))236 ax = plt.gca()237 for mask in outputs["masks"]:238 show_mask(mask, ax=ax, random_color=True)239 plt.axis("off")240 plt.show()241 if return_image:242 fig = plt.gcf()243 return fig2img(fig)244 245"""## imports"""246 247from transformers import pipeline248from transformers import SamModel, SamProcessor249from transformers import BlipForImageTextRetrieval250from transformers import AutoProcessor251 252from transformers.utils import logging253logging.set_verbosity_error()254#ignore_warnings()255 256import io257import matplotlib.pyplot as plt258import requests259import inflect260from PIL import Image261 262import os263import gradio as gr264 265import time266 267"""# Object detection268 269## hugging face model ("facebook/detr-resnet-50"). 167MB270"""271 272od_pipe = pipeline("object-detection", "facebook/detr-resnet-50")273 274chosen_model = pipeline("object-detection", "hustvl/yolos-small")275 276"""## gradio funcs"""277 278def get_object_detection_prediction(model_name, raw_image):279 model = od_pipe280 if "chosen-model" in model_name:281 model = chosen_model282 start = time.time()283 pipeline_output = model(raw_image)284 end = time.time()285 elapsed_result = f'{model_name} object detection elapsed {end-start} seconds'286 print(elapsed_result)287 processed_image = render_results_in_image(raw_image, pipeline_output)288 return [processed_image, elapsed_result]289 290"""# Image segmentation291 292## hugging face models: Zigeng/SlimSAM-uniform-77(segmentation) 39MB, Intel/dpt-hybrid-midas(depth) 490MB293"""294 295hugging_face_segmentation_pipe = pipeline("mask-generation", "Zigeng/SlimSAM-uniform-77")296hugging_face_segmentation_model = SamModel.from_pretrained("Zigeng/SlimSAM-uniform-77")297hugging_face_segmentation_processor = SamProcessor.from_pretrained("Zigeng/SlimSAM-uniform-77")298hugging_face_depth_estimator = pipeline(task="depth-estimation", model="Intel/dpt-hybrid-midas")299 300"""## chosen models: facebook/sam-vit-base(segmentation) 375MB, LiheYoung/depth-anything-small-hf(depth) 100MB"""301 302chosen_name = "facebook/sam-vit-base"303chosen_segmentation_pipe = pipeline("mask-generation", chosen_name)304chosen_segmentation_model = SamModel.from_pretrained(chosen_name)305chosen_segmentation_processor = SamProcessor.from_pretrained(chosen_name)306chosen_depth_estimator = pipeline(task="depth-estimation", model="LiheYoung/depth-anything-small-hf")307 308"""## gradio funcs"""309 310input_points = [[[1600, 700]]]311 312def segment_image_pretrained(model_name, raw_image):313 processor = hugging_face_segmentation_processor314 model = hugging_face_segmentation_model315 if("chosen" in model_name):316 processor = chosen_segmentation_processor317 model = chosen_segmentation_model318 start = time.time()319 inputs = processor(raw_image,320 input_points=input_points,321 return_tensors="pt")322 with torch.no_grad():323 outputs = model(**inputs)324 predicted_masks = processor.image_processor.post_process_masks(325 outputs.pred_masks,326 inputs["original_sizes"],327 inputs["reshaped_input_sizes"])328 results = []329 predicted_mask = predicted_masks[0]330 end = time.time()331 elapsed_result = f'{model_name} pretrained image segmentation elapsed {end-start} seconds'332 print(elapsed_result)333 for i in range(3):334 results.append(show_mask_on_image(raw_image, predicted_mask[:, i], return_image=True))335 results.append(elapsed_result);336 return results337 338def segment_image(model_name, raw_image):339 model = hugging_face_segmentation_pipe340 if("chosen" in model_name):341 print("chosen model used")342 model = chosen_segmentation_pipe343 start = time.time()344 output = model(raw_image, points_per_batch=32)345 end = time.time()346 elapsed_result = f'{model_name} raw image segmentation elapsed {end-start} seconds'347 print(elapsed_result)348 return [show_pipe_masks_on_image(raw_image, output, return_image = True), elapsed_result]349 350def depth_image(model_name, input_image):351 depth_estimator = hugging_face_depth_estimator352 print(model_name)353 if("chosen" in model_name):354 print("chosen model used")355 depth_estimator = chosen_depth_estimator356 start = time.time()357 out = depth_estimator(input_image)358 prediction = torch.nn.functional.interpolate(359 out["predicted_depth"].unsqueeze(0).unsqueeze(0),360 size=input_image.size[::-1],361 mode="bicubic",362 align_corners=False,363 )364 end = time.time()365 elapsed_result = f'{model_name} Depth Estimation elapsed {end-start} seconds'366 print(elapsed_result)367 output = prediction.squeeze().numpy()368 formatted = (output * 255 / np.max(output)).astype("uint8")369 depth = Image.fromarray(formatted)370 return [depth, elapsed_result]371 372"""# Image retrieval373 374## hugging face model: Salesforce/blip-itm-base-coco 900MB375"""376 377hugging_face_retrieval_model = BlipForImageTextRetrieval.from_pretrained(378 "Salesforce/blip-itm-base-coco")379hugging_face_retrieval_processor = AutoProcessor.from_pretrained(380 "Salesforce/blip-itm-base-coco")381 382"""## chosen model: Salesforce/blip-itm-base-flickr 900MB"""383 384chosen_retrieval_model = BlipForImageTextRetrieval.from_pretrained(385 "Salesforce/blip-itm-base-flickr")386chosen_retrieval_processor = AutoProcessor.from_pretrained(387 "Salesforce/blip-itm-base-flickr")388 389"""## gradion func"""390 391def retrieve_image(model_name, raw_image, predict_text):392 processor = hugging_face_retrieval_processor393 model = hugging_face_retrieval_model394 if("chosen" in model_name):395 processor = chosen_retrieval_processor396 model = chosen_retrieval_model397 start = time.time()398 inputs = processor(images=raw_image,399 text=predict_text,400 return_tensors="pt")401 end = time.time()402 elapsed_result = f"{model_name} image retrieval elapsed {end-start} seconds"403 print(elapsed_result)404 itm_scores = model(**inputs)[0]405 itm_score = torch.nn.functional.softmax(itm_scores,dim=1)406 return [f"""\407 The image and text are matched \408 with a probability of {itm_score[0][1]:.4f}""",409 elapsed_result]410 411"""# gradio"""412 413with gr.Blocks() as object_detection_tab:414 gr.Markdown("# Detect objects on image")415 gr.Markdown("Upload an image, choose model, press button.")416 417 with gr.Row():418 with gr.Column():419 # Input components420 input_image = gr.Image(label="Upload Image", type="pil")421 model_selector = gr.Dropdown(["hugging-face(facebook/detr-resnet-50)", "chosen-model(hustvl/yolos-small)"],422 label = "Select Model")423 424 with gr.Column():425 # Output image426 elapsed_result = gr.Textbox(label="Seconds elapsed", lines=1)427 output_image = gr.Image(label="Output Image", type="pil")428 429 # Process button430 process_btn = gr.Button("Detect objects")431 432 # Connect the input components to the processing function433 process_btn.click(434 fn=get_object_detection_prediction,435 inputs=[436 model_selector,437 input_image438 ],439 outputs=[output_image, elapsed_result]440 )441 442with gr.Blocks() as image_segmentation_detection_tab:443 gr.Markdown("# Image segmentation")444 gr.Markdown("Upload an image, choose model, press button.")445 446 with gr.Row():447 with gr.Column():448 # Input components449 input_image = gr.Image(label="Upload Image", type="pil")450 model_selector = gr.Dropdown(["hugging-face(Zigeng/SlimSAM-uniform-77)", "chosen-model(facebook/sam-vit-base)"],451 label = "Select Model")452 453 with gr.Column():454 elapsed_result = gr.Textbox(label="Seconds elapsed", lines=1)455 # Output image456 output_image = gr.Image(label="Segmented image", type="pil")457 with gr.Row():458 with gr.Column():459 segment_btn = gr.Button("Segment image(not pretrained)")460 461 with gr.Row():462 elapsed_result_pretrained_segment = gr.Textbox(label="Seconds elapsed", lines=1)463 with gr.Column():464 segment_pretrained_output_image_1 = gr.Image(label="Segmented image by pretrained model", type="pil")465 with gr.Column():466 segment_pretrained_output_image_2 = gr.Image(label="Segmented image by pretrained model", type="pil")467 with gr.Column():468 segment_pretrained_output_image_3 = gr.Image(label="Segmented image by pretrained model", type="pil")469 with gr.Row():470 with gr.Column():471 segment_pretrained_model_selector = gr.Dropdown(["hugging-face(Zigeng/SlimSAM-uniform-77)", "chosen-model(facebook/sam-vit-base)"],472 label = "Select Model")473 segment_pretrained_btn = gr.Button("Segment image(pretrained)")474 475 with gr.Row():476 with gr.Column():477 depth_output_image = gr.Image(label="Depth image", type="pil")478 elapsed_result_depth = gr.Textbox(label="Seconds elapsed", lines=1)479 with gr.Row():480 with gr.Column():481 depth_model_selector = gr.Dropdown(["hugging-face(Intel/dpt-hybrid-midas)", "chosen-model(LiheYoung/depth-anything-small-hf)"],482 label = "Select Model")483 depth_btn = gr.Button("Get image depth")484 485 segment_btn.click(486 fn=segment_image,487 inputs=[488 model_selector,489 input_image490 ],491 outputs=[output_image, elapsed_result]492 )493 segment_pretrained_btn.click(494 fn=segment_image_pretrained,495 inputs=[496 segment_pretrained_model_selector,497 input_image498 ],499 outputs=[segment_pretrained_output_image_1, segment_pretrained_output_image_2, segment_pretrained_output_image_3, elapsed_result_pretrained_segment]500 )501 502 depth_btn.click(503 fn=depth_image,504 inputs=[505 depth_model_selector,506 input_image,507 ],508 outputs=[depth_output_image, elapsed_result_depth]509 )510 511with gr.Blocks() as image_retrieval_tab:512 gr.Markdown("# Check is text describes image")513 gr.Markdown("Upload an image, choose model, press button.")514 515 with gr.Row():516 with gr.Column():517 # Input components518 input_image = gr.Image(label="Upload Image", type="pil")519 text_prediction = gr.TextArea(label="Describe image")520 model_selector = gr.Dropdown(["hugging-face(Salesforce/blip-itm-base-coco)", "chosen-model(Salesforce/blip-itm-base-flickr)"],521 label = "Select Model")522 523 with gr.Column():524 # Output image525 output_result = gr.Textbox(label="Probability result", lines=3)526 elapsed_result = gr.Textbox(label="Seconds elapsed", lines=1)527 528 # Process button529 process_btn = gr.Button("Detect objects")530 531 # Connect the input components to the processing function532 process_btn.click(533 fn=retrieve_image,534 inputs=[535 model_selector,536 input_image,537 text_prediction538 ],539 outputs=[output_result, elapsed_result]540 )541 542with gr.Blocks() as app:543 gr.TabbedInterface(544 [object_detection_tab,545 image_segmentation_detection_tab,546 image_retrieval_tab],547 ["Object detection",548 "Image segmentation",549 "Retrieve image"550 ],551 )552 553app.launch(share=True, debug=True)554 555app.close()