Basically I am trying to create a model that will detect angle on which a specific image is rotated. Also I have a dataset of 1500 documents resulting with images rotated on random.sample([0, 90, -90, 180], 2) and each of this angles has variation of random.uniform(-10, 10) resulting in ~4k rotated images. So I've come up with current model to predict sin and cos of the desired angle: class CnnRotateRegression(nn.Module): def __init__(self): super(CnnRotateRegression, self).__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=3) self.conv2 = nn.Conv2d(64, 128, kernel_size=3) self.conv3 = nn.Conv2d(128, 256, kernel_size=3) self.conv4 = nn.Conv2d(256, 512, kernel_size=3) self.conv5 = nn.Conv2d(512, 512, kernel_size=3) self.bn1 = nn.BatchNorm2d(64) self.bn2 = nn.BatchNorm2d(128) self.bn3 = nn.BatchNorm2d(256) self.bn4 = nn.BatchNorm2d(512) self.bn5 = nn.BatchNorm2d(512) self.activation = nn.ReLU() self.pool = nn.AvgPool2d(kernel_size=2) self.pool2 = nn.AdaptiveAvgPool2d((8,8)) self.linear_l1 = nn.Linear(512*8*8, 512) self.linear_l2 = nn.Linear(512, 256) self.linear_l3 = nn.Linear(256, 2) # sin + cos def forward(self, x): x = self.activation(self.pool(self.bn1(self.conv1(x)))) x = self.activation(self.pool(self.bn2(self.conv2(x)))) x = self.activation(self.pool(self.bn3(self.conv3(x)))) x = self.activation(self.pool(self.bn4(self.conv4(x)))) x = self.activation(self.pool(self.bn5(self.conv5(x)))) x = self.pool2(x) x = x.view(x.size(0), -1) x = self.activation(self.linear_l1(x)) x = self.activation(self.linear_l2(x)) x = self.linear_l3(x) x = F.normalize(x, p=2, dim=1) return x training part: model = CnnRotateRegression() model = model.to(device) loss_function = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=0.001) num_of_epochs = 11 for epoch in range(num_of_epochs): model.train() running_loss = 0.0 for images, labels in tqdm(train_Loader, desc="training loop"): images, labels = images.to(device), labels.to(device).float() angles = angle_to_sin_cos(labels) norm_angles = F.normalize(angles, p=2, dim=1) optimizer.zero_grad() outputs = model(images) loss = loss_function(outputs, norm_angles) loss.backward() optimizer.step() running_loss += loss.item() train_loss = running_loss / len(train_Loader) functions to convert sin and cos to angle and vice versa: def angle_to_sin_cos(angle): tensor_angle = angle.clone().detach() radian = tensor_angle * torch.pi / 180.0 return torch.stack([torch.cos(radian), torch.sin(radian)], dim=1) def sin_cos_to_angle(outputs): cos_val, sin_val = outputs[:, 0], outputs[:, 1] angle_rad = torch.atan2(sin_val, cos_val) angle_deg = angle_rad * (180 / torch.pi) return angle_deg My model performs poorly in determining small angles in range +-10 degrees. What would You suggest to improve/enhance to achieve better "small-degree" prediction? Thank you in advance! Continue reading...
