| 1101 | return boxes_filt, pred_phrases |
| 1102 | |
| 1103 | def plot_boxes_to_image(self, image_pil, tgt): |
| 1104 | H, W = tgt["size"] |
| 1105 | boxes = tgt["boxes"] |
| 1106 | labels = tgt["labels"] |
| 1107 | assert len(boxes) == len(labels), "boxes and labels must have same length" |
| 1108 | |
| 1109 | draw = ImageDraw.Draw(image_pil) |
| 1110 | mask = Image.new("L", image_pil.size, 0) |
| 1111 | mask_draw = ImageDraw.Draw(mask) |
| 1112 | |
| 1113 | # draw boxes and masks |
| 1114 | for box, label in zip(boxes, labels): |
| 1115 | # from 0..1 to 0..W, 0..H |
| 1116 | box = box * torch.Tensor([W, H, W, H]) |
| 1117 | # from xywh to xyxy |
| 1118 | box[:2] -= box[2:] / 2 |
| 1119 | box[2:] += box[:2] |
| 1120 | # random color |
| 1121 | color = tuple(np.random.randint(0, 255, size=3).tolist()) |
| 1122 | # draw |
| 1123 | x0, y0, x1, y1 = box |
| 1124 | x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1) |
| 1125 | |
| 1126 | draw.rectangle([x0, y0, x1, y1], outline=color, width=6) |
| 1127 | # draw.text((x0, y0), str(label), fill=color) |
| 1128 | |
| 1129 | font = ImageFont.load_default() |
| 1130 | if hasattr(font, "getbbox"): |
| 1131 | bbox = draw.textbbox((x0, y0), str(label), font) |
| 1132 | else: |
| 1133 | w, h = draw.textsize(str(label), font) |
| 1134 | bbox = (x0, y0, w + x0, y0 + h) |
| 1135 | # bbox = draw.textbbox((x0, y0), str(label)) |
| 1136 | draw.rectangle(bbox, fill=color) |
| 1137 | draw.text((x0, y0), str(label), fill="white") |
| 1138 | |
| 1139 | mask_draw.rectangle([x0, y0, x1, y1], fill=255, width=2) |
| 1140 | |
| 1141 | return image_pil, mask |
| 1142 | |
| 1143 | @prompts(name="Detect the Give Object", |
| 1144 | description="useful when you only want to detect or find out given objects in the picture" |