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:

https://preview.redd.it/ugq1rwkkhpt51.png?width=251&format=png&auto=webp&s=63c795850db1554f8be8d833fee95d3f96f3bc71

https://preview.redd.it/ua4tbszlhpt51.png?width=251&format=png&auto=webp&s=269a941f4747f64d5518013c57ff1234b4e6c549

https://preview.redd.it/qw70naonhpt51.png?width=251&format=png&auto=webp&s=5eb9a0b5adf7c7387e5492d817128f8335674efa

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]