chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Docker Image CI / build-ubuntu2004 (push) Waiting to run
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,255 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image as Image
|
||||
import PIL.ImageDraw as ImageDraw
|
||||
import PIL.ImageFont as ImageFont
|
||||
import PIL.ImageFilter as ImageFilter
|
||||
|
||||
|
||||
COLORS = [
|
||||
"GoldenRod",
|
||||
"MediumTurquoise",
|
||||
"GreenYellow",
|
||||
"SteelBlue",
|
||||
"DarkSeaGreen",
|
||||
"SeaShell",
|
||||
"LightGrey",
|
||||
"IndianRed",
|
||||
"DarkKhaki",
|
||||
"LawnGreen",
|
||||
"WhiteSmoke",
|
||||
"Peru",
|
||||
"LightCoral",
|
||||
"FireBrick",
|
||||
"OldLace",
|
||||
"LightBlue",
|
||||
"SlateGray",
|
||||
"OliveDrab",
|
||||
"NavajoWhite",
|
||||
"PaleVioletRed",
|
||||
"SpringGreen",
|
||||
"AliceBlue",
|
||||
"Violet",
|
||||
"DeepSkyBlue",
|
||||
"Red",
|
||||
"MediumVioletRed",
|
||||
"PaleTurquoise",
|
||||
"Tomato",
|
||||
"Azure",
|
||||
"Yellow",
|
||||
"Cornsilk",
|
||||
"Aquamarine",
|
||||
"CadetBlue",
|
||||
"CornflowerBlue",
|
||||
"DodgerBlue",
|
||||
"Olive",
|
||||
"Orchid",
|
||||
"LemonChiffon",
|
||||
"Sienna",
|
||||
"OrangeRed",
|
||||
"Orange",
|
||||
"DarkSalmon",
|
||||
"Magenta",
|
||||
"Wheat",
|
||||
"Lime",
|
||||
"GhostWhite",
|
||||
"SlateBlue",
|
||||
"Aqua",
|
||||
"MediumAquaMarine",
|
||||
"LightSlateGrey",
|
||||
"MediumSeaGreen",
|
||||
"SandyBrown",
|
||||
"YellowGreen",
|
||||
"Plum",
|
||||
"FloralWhite",
|
||||
"LightPink",
|
||||
"Thistle",
|
||||
"DarkViolet",
|
||||
"Pink",
|
||||
"Crimson",
|
||||
"Chocolate",
|
||||
"DarkGrey",
|
||||
"Ivory",
|
||||
"PaleGreen",
|
||||
"DarkGoldenRod",
|
||||
"LavenderBlush",
|
||||
"SlateGrey",
|
||||
"DeepPink",
|
||||
"Gold",
|
||||
"Cyan",
|
||||
"LightSteelBlue",
|
||||
"MediumPurple",
|
||||
"ForestGreen",
|
||||
"DarkOrange",
|
||||
"Tan",
|
||||
"Salmon",
|
||||
"PaleGoldenRod",
|
||||
"LightGreen",
|
||||
"LightSlateGray",
|
||||
"HoneyDew",
|
||||
"Fuchsia",
|
||||
"LightSeaGreen",
|
||||
"DarkOrchid",
|
||||
"Green",
|
||||
"Chartreuse",
|
||||
"LimeGreen",
|
||||
"AntiqueWhite",
|
||||
"Beige",
|
||||
"Gainsboro",
|
||||
"Bisque",
|
||||
"SaddleBrown",
|
||||
"Silver",
|
||||
"Lavender",
|
||||
"Teal",
|
||||
"LightCyan",
|
||||
"PapayaWhip",
|
||||
"Purple",
|
||||
"Coral",
|
||||
"BurlyWood",
|
||||
"LightGray",
|
||||
"Snow",
|
||||
"MistyRose",
|
||||
"PowderBlue",
|
||||
"DarkCyan",
|
||||
"White",
|
||||
"Turquoise",
|
||||
"MediumSlateBlue",
|
||||
"PeachPuff",
|
||||
"Moccasin",
|
||||
"LightSalmon",
|
||||
"SkyBlue",
|
||||
"Khaki",
|
||||
"MediumSpringGreen",
|
||||
"BlueViolet",
|
||||
"MintCream",
|
||||
"Linen",
|
||||
"SeaGreen",
|
||||
"HotPink",
|
||||
"LightYellow",
|
||||
"BlanchedAlmond",
|
||||
"RoyalBlue",
|
||||
"RosyBrown",
|
||||
"MediumOrchid",
|
||||
"DarkTurquoise",
|
||||
"LightGoldenRodYellow",
|
||||
"LightSkyBlue",
|
||||
]
|
||||
|
||||
|
||||
# Overlay mask with transparency on top of the image.
|
||||
def overlay(image, mask, color, alpha_transparency=0.5):
|
||||
for channel in range(3):
|
||||
image[:, :, channel] = np.where(
|
||||
mask == 1,
|
||||
image[:, :, channel] * (1 - alpha_transparency)
|
||||
+ alpha_transparency * color[channel] * 255,
|
||||
image[:, :, channel],
|
||||
)
|
||||
return image
|
||||
|
||||
|
||||
def visualize_detections(
|
||||
image_path, output_path, detections, labels=[], iou_threshold=0.5
|
||||
):
|
||||
image = Image.open(image_path).convert(mode="RGB")
|
||||
# Get image dimensions.
|
||||
im_width, im_height = image.size
|
||||
line_width = 2
|
||||
font = ImageFont.load_default()
|
||||
for d in detections:
|
||||
color = COLORS[d["class"] % len(COLORS)]
|
||||
# Dynamically convert PIL color into RGB numpy array.
|
||||
pixel_color = Image.new("RGB", (1, 1), color)
|
||||
# Normalize.
|
||||
np_color = (np.asarray(pixel_color)[0][0]) / 255
|
||||
# TRT instance segmentation masks.
|
||||
if isinstance(d["mask"], np.ndarray) and d["mask"].shape == (28, 28):
|
||||
# PyTorch uses [x1,y1,x2,y2] format instead of regular [y1,x1,y2,x2].
|
||||
d["ymin"], d["xmin"], d["ymax"], d["xmax"] = (
|
||||
d["xmin"],
|
||||
d["ymin"],
|
||||
d["xmax"],
|
||||
d["ymax"],
|
||||
)
|
||||
# Get detection bbox resolution.
|
||||
det_width = round(d["xmax"] - d["xmin"])
|
||||
det_height = round(d["ymax"] - d["ymin"])
|
||||
# Slight scaling, to get binary masks after float32 -> uint8
|
||||
# conversion, if not scaled all pixels are zero.
|
||||
mask = d["mask"] > iou_threshold
|
||||
# Convert float32 -> uint8.
|
||||
mask = mask.astype(np.uint8)
|
||||
# Create an image out of predicted mask array.
|
||||
small_mask = Image.fromarray(mask)
|
||||
# Upsample mask to detection bbox's size.
|
||||
mask = small_mask.resize((det_width, det_height), resample=Image.BILINEAR)
|
||||
# Create an original image sized template for correct mask placement.
|
||||
pad = Image.new("L", (im_width, im_height))
|
||||
# Place your mask according to detection bbox placement.
|
||||
pad.paste(mask, (round(d["xmin"]), (round(d["ymin"]))))
|
||||
# Reconvert mask into numpy array for evaluation.
|
||||
padded_mask = np.array(pad)
|
||||
# Creat np.array from original image, copy in order to modify.
|
||||
image_copy = np.asarray(image).copy()
|
||||
# Image with overlaid mask.
|
||||
masked_image = overlay(image_copy, padded_mask, np_color)
|
||||
# Reconvert back to PIL.
|
||||
image = Image.fromarray(masked_image)
|
||||
|
||||
# Bbox lines.
|
||||
draw = ImageDraw.Draw(image)
|
||||
draw.line(
|
||||
[
|
||||
(d["xmin"], d["ymin"]),
|
||||
(d["xmin"], d["ymax"]),
|
||||
(d["xmax"], d["ymax"]),
|
||||
(d["xmax"], d["ymin"]),
|
||||
(d["xmin"], d["ymin"]),
|
||||
],
|
||||
width=line_width,
|
||||
fill=color,
|
||||
)
|
||||
label = "Class {}".format(d["class"])
|
||||
if d["class"] < len(labels):
|
||||
label = "{}".format(labels[d["class"]])
|
||||
score = d["score"]
|
||||
text = "{}: {}%".format(label, int(100 * score))
|
||||
if score < 0:
|
||||
text = label
|
||||
left, top, right, bottom = font.getbbox(text)
|
||||
text_width, text_height = right - left, bottom - top
|
||||
text_bottom = max(text_height, d["ymin"])
|
||||
text_left = d["xmin"]
|
||||
margin = np.ceil(0.05 * text_height)
|
||||
draw.rectangle(
|
||||
[
|
||||
(text_left, text_bottom - text_height - 2 * margin),
|
||||
(text_left + text_width, text_bottom),
|
||||
],
|
||||
fill=color,
|
||||
)
|
||||
draw.text(
|
||||
(text_left + margin, text_bottom - text_height - margin),
|
||||
text,
|
||||
fill="black",
|
||||
font=font,
|
||||
)
|
||||
if output_path is None:
|
||||
return image
|
||||
image.save(output_path)
|
||||
Reference in New Issue
Block a user