0
$\begingroup$

I am dealing with the problem of estimating the angle of rotation of objects in images. The problem is that the network gets stuck when training at a loss level of about 90.

Below is the code for my loss function. This is a modified MAE.

class AngleLoss(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, input, target):
        diff=torch.abs(target-input)
        loss = torch.where(diff > 180, 360 - diff, diff)
        loss_m = loss.mean()
        return loss_m

During training, the date generator rotates the images by a random angle and then trims them to a circle of the diameter of the image size. I have checked several times if the generated data is definitely correct, but it does not seem to contain errors. Below is an example of the input image to the network and the data generation code.

enter image description here

class TrainDataLoader(Dataset):
    def __init__(self, data, size):
        self.data = data
        self.long = len(self.data)
        self.size = size
        self.transform_input_data = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
    
    def rotate_image(self, image, angle):
        angle = 360 - angle
        height, width = image.shape[:2]
        center = (width // 2, height // 2)
        border_value = (255, 255, 255)
        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        angle = np.deg2rad(angle)
        sin = np.sin(angle)
        cos = np.cos(angle)
        new_width = int((height * abs(sin)) + (width * abs(cos)))
        new_height = int((height * abs(cos)) + (width * abs(sin)))
        rotation_matrix[0, 2] += (new_width / 2) - center[0]
        rotation_matrix[1, 2] += (new_height / 2) - center[1]
        image = cv2.warpAffine(image, rotation_matrix, (new_width, new_height), borderValue=border_value)
        image = cv2.resize(image, (self.size, self.size))
        return image

    def circle_masking(self, image):
        h, w= image.shape[:2]
        c_x = w// 2
        c_y = h// 2
        mask = np.zeros_like(image)
        mask = cv2.circle(mask, (c_x, c_y), int(224/2), (255, 255, 255), -1)
        res = cv2.bitwise_and(image, mask)
        return res
    
    def __len__(self):
        return self.long
    
    def __getitem__(self, index):
        input = cv2.imread(self.data[index][0])
        input = cv2.resize(input, (224,224))
        target = self.data[index][1]
    
        angle = random.randint(0,359)

        input = self.rotate_image(input, angle)
        input = self.circle_masking(input)
    
        target += angle
        
        if target>360:
            target-=360

    
        input = self.transform_input_data(input)
        input[torch.isinf(input)] = 0
        input[torch.isnan(input)] = 0
        target = torch.tensor(target)

        return input, target

Finally, this is the code of my neural network and the plot of the loss function during training. enter image description here

class AngleTableNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V2)
        for param in self.model.parameters():
            param.requires_grad = True
        self.model.fc = nn.Linear(self.model.fc.in_features, 1)
    def forward(self, x):
        x = self.model(x)
        x = x.squeeze()
        x=x%360
        return x
$\endgroup$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.