| 50 | return image |
| 51 | |
| 52 | def visualize_detections(image_path, output_path, detections, labels=[], iou_threshold=0.5): |
| 53 | image = Image.open(image_path).convert(mode='RGB') |
| 54 | # Get image dimensions. |
| 55 | im_width, im_height = image.size |
| 56 | line_width = 2 |
| 57 | font = ImageFont.load_default() |
| 58 | for d in detections: |
| 59 | color = COLORS[d['class'] % len(COLORS)] |
| 60 | # Dynamically convert PIL color into RGB numpy array. |
| 61 | pixel_color = Image.new("RGB",(1, 1), color) |
| 62 | # Normalize. |
| 63 | np_color = (np.asarray(pixel_color)[0][0])/255 |
| 64 | # TRT instance segmentation masks. |
| 65 | if isinstance(d['mask'], np.ndarray) and d['mask'].shape == (28, 28): |
| 66 | # PyTorch uses [x1,y1,x2,y2] format instead of regular [y1,x1,y2,x2]. |
| 67 | d['ymin'], d['xmin'], d['ymax'], d['xmax'] = d['xmin'], d['ymin'], d['xmax'], d['ymax'] |
| 68 | # Get detection bbox resolution. |
| 69 | det_width = round(d['xmax'] - d['xmin']) |
| 70 | det_height = round(d['ymax'] - d['ymin']) |
| 71 | # Slight scaling, to get binary masks after float32 -> uint8 |
| 72 | # conversion, if not scaled all pixels are zero. |
| 73 | mask = d['mask'] > iou_threshold |
| 74 | # Convert float32 -> uint8. |
| 75 | mask = mask.astype(np.uint8) |
| 76 | # Create an image out of predicted mask array. |
| 77 | small_mask = Image.fromarray(mask) |
| 78 | # Upsample mask to detection bbox's size. |
| 79 | mask = small_mask.resize((det_width, det_height), resample=Image.BILINEAR) |
| 80 | # Create an original image sized template for correct mask placement. |
| 81 | pad = Image.new("L", (im_width, im_height)) |
| 82 | # Place your mask according to detection bbox placement. |
| 83 | pad.paste(mask, (round(d['xmin']), (round(d['ymin'])))) |
| 84 | # Reconvert mask into numpy array for evaluation. |
| 85 | padded_mask = np.array(pad) |
| 86 | #Creat np.array from original image, copy in order to modify. |
| 87 | image_copy = np.asarray(image).copy() |
| 88 | # Image with overlaid mask. |
| 89 | masked_image = overlay(image_copy, padded_mask, np_color) |
| 90 | # Reconvert back to PIL. |
| 91 | image = Image.fromarray(masked_image) |
| 92 | |
| 93 | # Bbox lines. |
| 94 | draw = ImageDraw.Draw(image) |
| 95 | draw.line([(d['xmin'], d['ymin']), (d['xmin'], d['ymax']), (d['xmax'], d['ymax']), (d['xmax'], d['ymin']), |
| 96 | (d['xmin'], d['ymin'])], width=line_width, fill=color) |
| 97 | label = "Class {}".format(d['class']) |
| 98 | if d['class'] < len(labels): |
| 99 | label = "{}".format(labels[d['class']]) |
| 100 | score = d['score'] |
| 101 | text = "{}: {}%".format(label, int(100 * score)) |
| 102 | if score < 0: |
| 103 | text = label |
| 104 | text_width, text_height = font.getsize(text) |
| 105 | text_bottom = max(text_height, d['ymin']) |
| 106 | text_left = d['xmin'] |
| 107 | margin = np.ceil(0.05 * text_height) |
| 108 | draw.rectangle([(text_left, text_bottom - text_height - 2 * margin), (text_left + text_width, text_bottom)], |
| 109 | fill=color) |