Upload inference.py with huggingface_hub

inference.py

@@ -59,9 +59,36 @@ def resize_and_pad_image(image, new_shape, stride=32):
     return image
 
 
+def scale_boxes(img1_shape, boxes, img0_shape):
+    """
+    Rescales bounding boxes (in the format of xyxy by default) from the shape of the image they were originally
+    specified in (img1_shape) to the shape of a different image (img0_shape).
+
+    Args:
+        img1_shape (tuple): The shape of the image that the bounding boxes are for, in the format of (height, width).
+        boxes (torch.Tensor): the bounding boxes of the objects in the image, in the format of (x1, y1, x2, y2)
+        img0_shape (tuple): the shape of the target image, in the format of (height, width).
+
+    Returns:
+        boxes (torch.Tensor): The scaled bounding boxes, in the format of (x1, y1, x2, y2)
+    """
+
+    # Calculate scaling ratio
+    gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])
+
+    # Calculate padding size
+    pad_x = round((img1_shape[1] - img0_shape[1] * gain) / 2 - 0.1)
+    pad_y = round((img1_shape[0] - img0_shape[0] * gain) / 2 - 0.1)
+
+    # Remove padding and scale boxes
+    boxes[..., :4] = (boxes[..., :4] - [pad_x, pad_y, pad_x, pad_y]) / gain
+    return boxes
+
+
 class YoloResult:
     def __init__(self, boxes, names):
         self.boxes = [YoloBox(data=d) for d in boxes]
+        self.boxes = sorted(self.boxes, key=lambda x: x.conf, reverse=True)
         self.names = names
 
 
@@ -84,23 +111,28 @@ def inference(image):
     """
 
     # Preprocess image
+    orig_h, orig_w = image.shape[:2]
     image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
     pix = resize_and_pad_image(image, new_shape=int(image.shape[0] / stride) * stride)
     pix = np.transpose(pix, (2, 0, 1))  # CHW
     pix = np.expand_dims(pix, axis=0)  # BCHW
     pix = pix.astype(np.float32) / 255.0  # Normalize to [0, 1]
+    new_h, new_w = pix.shape[2:]
 
     # Run inference
     preds = session.run(None, {"images": pix})[0]
 
     # Postprocess predictions
     preds = preds[preds[..., 4] > 0.25]
+    preds[..., :4] = scale_boxes((new_h, new_w), preds[..., :4], (orig_h, orig_w))
     return YoloResult(boxes=preds, names=names)
 
 
 if __name__ == "__main__":
     import sys
+    import matplotlib
     import matplotlib.pyplot as plt
+    import matplotlib.colors as colors
 
     image = sys.argv[1]
     image = cv2.imread(image)
@@ -125,7 +157,18 @@ if __name__ == "__main__":
             bitmap[y0:y1, x0:x1] = i + 2
     bitmap = bitmap[::-1, :]
 
-    fig, ax = plt.subplots(1, 2, figsize=(10, 6))
+    # map bitmap to color
+    colormap = matplotlib.colormaps["Pastel1"]
+    norm = colors.Normalize(vmin=bitmap.min(), vmax=bitmap.max())
+    colored_bitmap = colormap(norm(bitmap))
+    colored_bitmap = (colored_bitmap[:, :, :3] * 255).astype(np.uint8)
+
+    # overlay bitmap on image
+    image_with_bitmap = cv2.multiply(image, colored_bitmap, scale=1 / 255)
+
+    # show the results
+    fig, ax = plt.subplots(1, 3, figsize=(15, 6))
     ax[0].imshow(image)
-    ax[1].imshow(bitmap)
+    ax[1].imshow(bitmap, cmap="Pastel1")
+    ax[2].imshow(image_with_bitmap)
     plt.show()