I am working on an algorithm for detecting keys and keyboard body from an image. The model of the keyboard is known, so a blender env. has been created to generate images with random lighting, angle, textures and objects on screen for training. I have chosen an approach utilising a UNet with the following structure:
class RELUConvBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() layers = [ nn.Conv2d(in_ch, out_ch,3,1,1), nn.BatchNorm2d(out_ch), nn.ReLU() ] self.model = nn.Sequential(*layers) def forward(self,x): return self.model(x) class DownBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() layers = [ RELUConvBlock(in_ch, out_ch), RELUConvBlock(out_ch, out_ch) ] self.model = nn.Sequential(*layers) def forward(self,x): return self.model(x) class UpBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() layers = [ RELUConvBlock(in_ch, out_ch), RELUConvBlock(out_ch, out_ch) ] self.model = nn.Sequential(*layers) def forward(self,x): return self.model(x) class UNet(nn.Module): def __init__(self, out_ch = 3, down_ch = [64,128,256,512]): super().__init__() self.pool = nn.MaxPool2d(kernel_size=(2,2), stride=(2,2)) self.down0 = DownBlock(3, down_ch[0]) self.down1 = DownBlock(down_ch[0], down_ch[1]) self.down2 = DownBlock(down_ch[1], down_ch[2]) self.down3 = DownBlock(down_ch[2], down_ch[3]) self.bottleneck = DownBlock(down_ch[3], 2*down_ch[3]) self.up3 = UpBlock(2*down_ch[-1], down_ch[-1]) self.up2 = UpBlock(down_ch[-1], down_ch[-2]) self.up1 = UpBlock(down_ch[-2], down_ch[-3]) self.up0 = UpBlock(down_ch[-3], down_ch[-4]) self.connect_b_up3 = nn.ConvTranspose2d(2*down_ch[-1], down_ch[-1],kernel_size=2,stride=2) self.connect_up3_up2 = nn.ConvTranspose2d(down_ch[-1], down_ch[-2],kernel_size=2,stride=2) self.connect_up2_up1 = nn.ConvTranspose2d(down_ch[-2], down_ch[-3],kernel_size=2,stride=2) self.connect_up1_up0 = nn.ConvTranspose2d(down_ch[-3], down_ch[-4],kernel_size=2,stride=2) self.final_conv = nn.Conv2d(down_ch[0], out_ch, kernel_size=1) def _crop_to_match(self, tensor, target): _, _, h, w = target.shape return tensor[:, :, :h, :w] def forward(self, x): skip_connections = [] x = self.down0(x) skip_connections.append(x) x = self.pool(x) x = self.down1(x) skip_connections.append(x) x = self.pool(x) x = self.down2(x) skip_connections.append(x) x = self.pool(x) x = self.down3(x) skip_connections.append(x) x = self.pool(x) x = self.bottleneck(x) x = self.connect_b_up3(x) skip_connection = self._crop_to_match(skip_connections[3], x) x = torch.cat((skip_connection, x), dim=1) x = self.up3(x) x = self.connect_up3_up2(x) skip_connection = self._crop_to_match(skip_connections[2], x) x = torch.cat((skip_connection, x), dim=1) x = self.up2(x) x = self.connect_up2_up1(x) skip_connection = self._crop_to_match(skip_connections[1], x) x = torch.cat((skip_connection, x), dim=1) x = self.up1(x) x = self.connect_up1_up0(x) skip_connection = self._crop_to_match(skip_connections[0], x) x = torch.cat((skip_connection, x), dim=1) x = self.up0(x) x = self.final_conv(x) return x #nn.functional.sigmoid(x) The network is trained on 500 images generated from the blender env.
LEARNING_RATE = 1e-4 num_epochs = 10 loss_fn = nn.CrossEntropyLoss() optimizer = Adam(model.parameters(), lr=LEARNING_RATE) scaler = torch.amp.GradScaler(device) model.train() for epoch in range(num_epochs): loop = tqdm(enumerate(train_loader), total = len(train_loader)) for batch_idx, (data, targets) in loop: data = data.to(device) targets = targets.to(device) with torch.amp.autocast(device_type=str(device)): predictions = model(data) loss = loss_fn(predictions, targets) optimizer.zero_grad() scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() loop.set_postfix(loss = loss.item()) While the results on training and testing data are very promising (results_1), when tested on real photos the network performs poorly (results_2). What can be done to fix that? Is there any way other than creating a better blender env? Maybe a different type of nn or a different solution entirely? This is my first time using AI to solve real world problem so i probably made a lot of bad choices.
I have tried to tweak the parameters, make the blender images darker to reflect reality better and using openCV to manipulate with the results but none of that worked well
