<jupyter_start><jupyter_text>covariant.ai Research Take-home: Oriented Bounding BoxesDependencies: `pip install -r requirements.txt` before beginning.We have also provided functions for visualizing the data and evaluating results (`visualize()` and `benchmark()`).<jupyter_code>from google.colab import files

uploaded = files.upload()

for fn in uploaded.keys():
  print('User uploaded file "{name}" with length {length} bytes'.format(
      name=fn, length=len(uploaded[fn])))
from collections import defaultdict
import numpy as np
import sys
import tqdm

import shapely.geometry as shg

from matplotlib import pyplot as plt
%matplotlib inline

def visualize(image, boxes=[], colors='rmcyg'):
    """
    Visualize an image and (optionally) bounding boxes.
    
    Args:
      image: numpy array of shape(height, width).
      boxes: a list or numpy array of bounding boxes, shape (N, 4, 2)
      colors:
    """
    plt.imshow(image)
    for i, b in enumerate(boxes):
        plt.plot(b[(0, 1, 2, 3, 0), 0], b[(0, 1, 2, 3, 0), 1], c=colors[i % len(colors)])

def benchmark(y_pred, y_true):
    """
    Compute iou among a batch of predictions.
    
    Args:
      y_pred: numpy array of shape (N,4,2)
      y_true: numpy array of shape (N,4,2)
      
    Returns:
      result: numpy array of shape(N), where `result[i] = iou(y_pred[i], y_true[i])`.
    """
    result = []
    for yp, yt in tqdm.tqdm(zip(y_pred, y_true)):
        a, b = shg.Polygon(yt), shg.Polygon(yp)
                
        # Avoid issues with ill-formed polygons that have self intersection
        a = a.buffer(0)
        b = b.buffer(0)
        
        result.append(a.intersection(b).area / a.union(b).area)
    return np.array(result)<jupyter_output><empty_output><jupyter_text>Your Solution Here: Observations:1. digits are of different scales/orientation2. different orientation.Step 1: turn left, right, top, bottom into center, height, width, angleStep 2: building a vanilla cnn as baseline to learn the oriented box.Step 3: add some tricks to learn the oriented *box*. orientation go from 0 to 180. use tan theta<jupyter_code># import matplotlib.pyplot as plt
# range of input. int, 255. 0 is dominated. --> gray scale to binary? probably easier.
# plt.hist(test_dset["images"].flatten(), bins = list(range(2, 255, 8)))
# plt.hist(test_dset["images"].flatten(), bins = list(range(0, 255, 8)))
# prepare dataset.
import torchvision
import torch
from torchvision.transforms import ToTensor
from torch.utils.data import Dataset
class MINISTDataset(Dataset):
    def __init__(self, path):
        self.path = path
        self.dset = dict(np.load(path))

    def __len__(self):
        return self.dset["images"].shape[0]

    def __getitem__(self, idx):
        # get item
        image = self.dset["images"][idx, :, :]
        image = np.expand_dims(image, axis=-1)
        box = self.dset["boxes"][idx, :, :]

        # transform image: normalize to [-1, 1]
        image = ToTensor()(image)
        # transform box
        label = torch.tensor(box.flatten(), dtype=torch.float32)

        return image, label



def visualize_minist_dataset(dataset, plot_predictions=False, predictions=None):
    def plot_box(label, color = 'g'):
        b = label.reshape(4, 2).numpy()
        plt.plot(b[(0, 1, 2, 3, 0), 0], b[(0, 1, 2, 3, 0), 1], c=color)
        plt.scatter(b[0, 0], b[0, 1], c='r', marker='*')
        # plot center
        center = (b[0, :] + b[2,:])/2
        plt.scatter(center[0], center[1], c='r', marker='*')

        # plot H, and W:
        h1= np.linalg.norm(b[0, :]- b[1, :])
        h2= np.linalg.norm(b[1, :]-b[2, :])
        plt.title(f"H {h1}, \n W {h2}")

    figure = plt.figure(figsize=(8, 8))
    cols, rows = 3, 3
    for i in range(1, cols * rows + 1):
        sample_idx = i
        img, label = dataset[sample_idx]
        figure.add_subplot(rows, cols, i)

        plt.axis("off")
        plt.imshow(img.squeeze(), cmap="gray")
        # plt.imshow(image)
        plot_box(label)
       
        if plot_predictions:
          prediction = predictions[sample_idx, :]
          plot_box(prediction, color='r')
          print(label)
          print(prediction)

    plt.show()


batch_size = 10

# iterating and visualizing the dataset
train_dataset = MINISTDataset("train.npz")
test_dataset = MINISTDataset("test.npz")

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# template = np.array([[-1, 1, 1, -1], [1, 1, -1, -1], [1, 1, 1, 1]])
# A_forward =  np.array([[hx * np.cos(theta), hy * -np.sin(theta), cx], 
#                       [hx * np.sin(theta), hy * np.cos(theta), cy]])
# rotated = np.transpose(np.matmul(A_forward, np.transpose(original)))



# use optimization algorithm to find the fit. interesting.
import numpy as np
from scipy.optimize import minimize
# def box_cxcywh_to_xyxy(x):
#     x_c, y_c, w, h = x.unbind(-1)
#     b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
#          (x_c + 0.5 * w), (y_c + 0.5 * h)]
#     return torch.stack(b, dim=-1)

def rotate_forward(hx, hy, cx, cy, theta, points):
  # points = [2, N] 4 points.
  rotation_matrix = np.array([[np.cos(theta), np.sin(theta)], 
                      [-np.sin(theta), np.cos(theta)]]).squeeze()
  # first, scale
  x = np.multiply(points, np.expand_dims(np.array([hx, hy]), axis=1))
  # second, rotate
  x = np.matmul(rotation_matrix, x)
  # third, shift
  x = x + np.expand_dims(np.array([cx, cy]), axis=1)
  return x

def rotate_backward(hx, hy, cx, cy, theta, points):
  # points = [2, N] 4 points.

  rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)], 
                      [np.sin(theta), np.cos(theta)]])
  # first, shift back
  x = points - np.expand_dims(np.array([cx, cy]), axis=1)
  # second, rotate back
  x = np.matmul(rotation_matrix, x)
  # third, scale back
  x = np.multiply(x, np.expand_dims(np.array([1./hx, 1./hy]), axis=1))
  return x

def construct_minf(hx, hy, cx, cy, source, target):
  def fmin(theta):
    return np.square(rotate_forward(hx, hy, cx, cy, theta, source) - target).sum()
  return fmin

def box_xyxy_to_cxcyhxhytheta(box):
  # b [2, 4]
   # center
    center = (box[:, 0] + box[:, 2])/2
    cx, cy = center[0],  center[1]
    # H and W
    hx= np.linalg.norm(box[:, 1]-box[:, 2])/2
    hy= np.linalg.norm(box[:, 0]-box[:, 1])/2
    
    # guess theta:
    source = np.transpose(np.array([[-1, -1], [-1, 1], [1, 1], [1, -1]])) # [2, 4]
    target = box
    fmin = construct_minf(hx, hy, cx, cy, source, target)
    x0 = np.array([np.pi])
    res = minimize(fmin, x0, method='nelder-mead', options={'xatol': 1e-8, 'disp': True})
    theta = res.x[0]
    
    return hx, hy, cx, cy, theta

def plot_box(b, color = 'g'):
  plt.plot(b[0,(0, 1, 2, 3, 0)], b[1, (0, 1, 2, 3, 0)], c=color)
  plt.scatter(b[0, 0], b[1, 0], c=color, marker='*', linewidths=3)
  center = (b[:, 0] + b[:, 2])/2 # [2, 1]
  plt.scatter(center[0], center[1], c=color, marker='*')

def visualize_rotation(img=None, label=None, hx=10, hy=5, cx=20, cy=3, theta = np.pi/3):
  if label is not None:
    box = np.transpose(label.reshape(4, 2).numpy()) # [2, 4]
    hx, hy, cx, cy, theta = box_xyxy_to_cxcyhxhytheta(box)
  source = np.transpose(np.array([[-1, -1], [-1, 1], [1, 1], [1, -1]])) # [2, 4]
  rotated = rotate_forward(hx, hy, cx, cy, theta, source)
  original_recover = rotate_backward(hx, hy, cx, cy, theta, rotated)

  # visualize it. plot box on top of images...
  figure = plt.figure(figsize=(8, 8))
  if img is not None:
    plt.imshow(img.squeeze(), cmap="gray")
  if label is not None:
    plot_box(box, color = 'g')
  plt.axis("equal")
  plot_box(rotated,  color='y')
  plot_box(original_recover,  color='r')
  plt.show()

img, label = test_dataset[8]
visualize_rotation(img, label=label)
visualize_minist_dataset(test_dataset)
import torch 
import torch.nn as nn
import torch.nn.functional as F
from torchvision.ops.boxes import box_area

def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(-1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=-1)


def box_xyxy_to_cxcywh(x):
    x0, y0, x1, y1 = x.unbind(-1)
    b = [(x0 + x1) / 2, (y0 + y1) / 2,
         (x1 - x0), (y1 - y0)]
    return torch.stack(b, dim=-1)


# modified from torchvision to also return the union
def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    iou = inter / union
    return iou, union


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        # Spatial transformer localization-network
        self.localization = nn.Sequential(
            nn.Conv2d(1, 8, kernel_size=7),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(8, 10, kernel_size=5),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True)
        )

        # Regressor for the hx, hy, cx, cy, theta
        self.fc_loc = nn.Sequential(
            nn.Linear(10 * 4 * 4, 32),
            nn.ReLU(True),
            nn.Linear(32, 5)
        )

        # Initialize the weights/bias with identity transformation
        self.fc_loc[2].weight.data.zero_()
        self.fc_loc[2].bias.data.copy_(torch.tensor([0, 0, 16, 16, 0], dtype=torch.float))

        # template
        self.template = torch.tensor([[-1, 1, 1],[1, -1, 1]], dtype=torch.float32).permute(1, 0).unsqueeze(0)
        print(self.template)
        print(self.template.shape)

    # Spatial transformer network forward function
    def stn(self, x):
        xs = self.localization(x)
        print(xs.shape)
        xs = xs.view(-1, 10 * 4 * 4)
        print(xs.shape)
        theta = self.fc_loc(xs)
        theta = theta.view(-1, 2, 3)
        print(theta.shape)
        print(self.template.shape)
        box = torch.matmul(theta, self.template)
        return box

    def forward(self, x, y):
        # transform the input
        x = self.stn(x)
        return x


model = Net()
import torch.optim as optim

model = Model(num_classes=1)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
criterion = nn.MSELoss()
# train the network
n_epoch = 2
for epoch in range(n_epoch):
      running_loss = 0.0
      for i, data in enumerate(train_loader, 0):
  # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data
#         labels = labels.reshape(-1, 4, 2).numpy()

#  b = label.reshape(4, 2).numpy()
#         plt.plot(b[(0, 1, 2, 3, 0), 0], b[(0, 1, 2, 3, 0), 1], c=color)
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % 20 == 19:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0


print('Finished Training')
PATH = './cifar_net.pth'
torch.save(model.state_dict(), PATH)

visualize_minist_dataset(test_dataset, colors='rmcyg')

model_load = Model()
model_load.load_state_dict(torch.load(PATH))
# evaluate all test data.

error = 0
total = 0
ntest = len(test_loader) * batch_size
print(ntest)
predictions = torch.zeros((ntest, 8))
print(predictions.shape)
# since we're not training, we don't need to calculate the gradients for our outputs
with torch.no_grad():
    for i, data in enumerate(test_loader, 0):
        images, labels = data
        # calculate outputs by running images through the network
        predicted = model_load(images)
        predictions[i*batch_size:min((i + 1)*batch_size, ntest), :] = predicted
        total += labels.size(0)
        error += ((predicted - labels).square()).sum().item()

print(f'Accuracy of the network on the 10000 test images: {error/total}')
# outputs = Model(images)
visualize_minist_dataset(test_dataset, plot_predictions=True, predictions=predictions)

i = 0
model = Net()
for data in train_loader:
  input, labels = data 
  output = model(input)
  print(output.shape)
  i += 1
  if i > 1:
    break<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15751295336787566, 'fraction_numerical': 0.02821841371064169, 'mean_word_length': 2.855562384757222, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '477741', 'license': 'MIT', 'n_tokens_mistral': 4840, 'n_tokens_neox': 4566, 'n_words': 1199, 'path': 'oriented_bbox_v2.ipynb', 'repo_name': 'Juyue/TextRecognitionDataGenerator', 'size': '12545'}