Initial commit

ML/Pytorch/Basics/albumentations_tutorial/utils.py

@@ -0,0 +1,36 @@
+import random
+import cv2
+from matplotlib import pyplot as plt
+import matplotlib.patches as patches
+import numpy as np
+import albumentations as A
+
+
+def visualize(image):
+    plt.figure(figsize=(10, 10))
+    plt.axis('off')
+    plt.imshow(image)
+    plt.show()
+
+
+def plot_examples(images, bboxes=None):
+    fig = plt.figure(figsize=(15, 15))
+    columns = 4
+    rows = 5
+
+    for i in range(1, len(images)):
+        if bboxes is not None:
+            img = visualize_bbox(images[i - 1], bboxes[i - 1], class_name="Elon")
+        else:
+            img = images[i-1]
+        fig.add_subplot(rows, columns, i)
+        plt.imshow(img)
+    plt.show()
+
+
+# From https://albumentations.ai/docs/examples/example_bboxes/
+def visualize_bbox(img, bbox, class_name, color=(255, 0, 0), thickness=5):
+    """Visualizes a single bounding box on the image"""
+    x_min, y_min, x_max, y_max = map(int, bbox)
+    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color, thickness)
+    return img