How to dettect rectangles in a simple image [Pytorch]

Original article was published by /u/tejonaco on Deep Learning

Hi, I am playing with neural networks trying to solve a simple problem, detect black rectangles in a white background. This is my model so far:

class RectangleDetector(nn.Module): def __init__(self): super(SquareDetector, self).__init__() ##LAYERS## self.conv1 = nn.Conv2d(in_channels=1, out_channels=4, kernel_size=3) self.l1 = nn.Linear(in_features=38416, out_features=4) def forward(self, x): x = self.conv1(x) x = x.flatten() x = self.l1(x) x = F.relu(x) return x 

This is a simple model for a simple task, I tried other more complex but it seems that the result is poorest. This are some results, the model knows something about the task but the result is not good:

The gray border is the model result.

A training with 1000-10000 samples give the best result (like this) so I think the model is not well configured and is not an poor training problem.

Thanks for your help.

submitted by /u/tejonaco
[link] [comments]