Spaces:
Runtime error
Runtime error
| import cv2 | |
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import logging | |
| from utils_model import compute_boxes_and_sizes, get_upsample_output, get_box_and_dot_maps, get_boxed_img | |
| from time import time | |
| class LSCCNN(nn.Module): | |
| def __init__(self, name='scale_4', checkpoint_path=None, output_downscale=2, | |
| PRED_DOWNSCALE_FACTORS=(8, 4, 2, 1), GAMMA=(1, 1, 2, 4), NUM_BOXES_PER_SCALE=3): | |
| super(LSCCNN, self).__init__() | |
| self.name = name | |
| if torch.cuda.is_available(): | |
| self.rgb_means = torch.cuda.FloatTensor([104.008, 116.669, 122.675]) | |
| else: | |
| self.rgb_means = torch.FloatTensor([104.008, 116.669, 122.675]) | |
| self.rgb_means = torch.autograd.Variable(self.rgb_means, requires_grad=False).unsqueeze(0).unsqueeze( | |
| 2).unsqueeze(3) | |
| self.BOXES, self.BOX_SIZE_BINS = compute_boxes_and_sizes(PRED_DOWNSCALE_FACTORS, GAMMA, NUM_BOXES_PER_SCALE) | |
| self.output_downscale = output_downscale | |
| in_channels = 3 | |
| self.relu = nn.ReLU(inplace=True) | |
| self.conv1_1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1) | |
| self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1) | |
| self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) | |
| self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1) | |
| self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1) | |
| self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) | |
| self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1) | |
| self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1) | |
| self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1) | |
| self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) | |
| self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1) | |
| self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1) | |
| self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1) | |
| self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2) | |
| self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1) | |
| self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1) | |
| self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1) | |
| self.convA_1 = nn.Conv2d(256, 256, kernel_size=3, padding=1) | |
| self.convA_2 = nn.Conv2d(256, 128, kernel_size=3, padding=1) | |
| self.convA_3 = nn.Conv2d(128, 64, kernel_size=3, padding=1) | |
| self.convA_4 = nn.Conv2d(64, 32, kernel_size=3, padding=1) | |
| self.convA_5 = nn.Conv2d(32, 4, kernel_size=3, padding=1) | |
| self.convB_1 = nn.Conv2d(512, 256, kernel_size=3, padding=1) | |
| self.convB_2 = nn.Conv2d(256, 128, kernel_size=3, padding=1) | |
| self.convB_3 = nn.Conv2d(128, 64, kernel_size=3, padding=1) | |
| self.convB_4 = nn.Conv2d(64, 32, kernel_size=3, padding=1) | |
| self.convB_5 = nn.Conv2d(32, 4, kernel_size=3, padding=1) | |
| self.convC_1 = nn.Conv2d(384, 256, kernel_size=3, padding=1) | |
| self.convC_2 = nn.Conv2d(256, 128, kernel_size=3, padding=1) | |
| self.convC_3 = nn.Conv2d(128, 64, kernel_size=3, padding=1) | |
| self.convC_4 = nn.Conv2d(64, 32, kernel_size=3, padding=1) | |
| self.convC_5 = nn.Conv2d(32, 4, kernel_size=3, padding=1) | |
| self.convD_1 = nn.Conv2d(256, 256, kernel_size=3, padding=1) | |
| self.convD_2 = nn.Conv2d(256, 128, kernel_size=3, padding=1) | |
| self.convD_3 = nn.Conv2d(128, 64, kernel_size=3, padding=1) | |
| self.convD_4 = nn.Conv2d(64, 32, kernel_size=3, padding=1) | |
| self.convD_5 = nn.Conv2d(32, 4, kernel_size=3, padding=1) | |
| self.conv_before_transpose_1 = nn.Conv2d(512, 256, kernel_size=3, padding=1) | |
| self.transpose_1 = nn.ConvTranspose2d(256, 256, kernel_size=3, stride=2, padding=1, output_padding=1) | |
| self.conv_after_transpose_1_1 = nn.Conv2d(256, 256, kernel_size=3, padding=1) | |
| self.transpose_2 = nn.ConvTranspose2d(256, 256, kernel_size=3, stride=2, padding=1, output_padding=1) | |
| self.conv_after_transpose_2_1 = nn.Conv2d(256, 128, kernel_size=3, padding=1) | |
| self.transpose_3 = nn.ConvTranspose2d(256, 256, kernel_size=3, stride=4, padding=0, output_padding=1) | |
| self.conv_after_transpose_3_1 = nn.Conv2d(256, 128, kernel_size=3, padding=1) | |
| self.transpose_4_1_a = nn.ConvTranspose2d(256, 256, kernel_size=3, stride=4, padding=0, output_padding=1) | |
| self.transpose_4_1_b = nn.ConvTranspose2d(256, 256, kernel_size=3, stride=2, padding=1, output_padding=1) | |
| self.conv_after_transpose_4_1 = nn.Conv2d(256, 64, kernel_size=3, padding=1) | |
| self.transpose_4_2 = nn.ConvTranspose2d(256, 256, kernel_size=3, stride=4, padding=0, output_padding=1) | |
| self.conv_after_transpose_4_2 = nn.Conv2d(256, 64, kernel_size=3, padding=1) | |
| self.transpose_4_3 = nn.ConvTranspose2d(128, 128, kernel_size=3, stride=2, padding=1, output_padding=1) | |
| self.conv_after_transpose_4_3 = nn.Conv2d(128, 64, kernel_size=3, padding=1) | |
| self.conv_middle_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1) | |
| self.conv_middle_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1) | |
| self.conv_middle_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1) | |
| self.conv_mid_4 = nn.Conv2d(512, 256, kernel_size=3, padding=1) | |
| self.conv_lowest_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1) | |
| self.conv_lowest_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1) | |
| self.conv_lowest_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1) | |
| self.conv_lowest_4 = nn.Conv2d(256, 128, kernel_size=3, padding=1) | |
| self.conv_scale1_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1) | |
| self.conv_scale1_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1) | |
| self.conv_scale1_3 = nn.Conv2d(128, 64, kernel_size=3, padding=1) | |
| if checkpoint_path is not None: | |
| self.load_state_dict(torch.load(checkpoint_path)) | |
| def forward(self, x): | |
| mean_sub_input = x | |
| mean_sub_input -= self.rgb_means | |
| #################### Stage 1 ########################## | |
| main_out_block1 = self.relu(self.conv1_2(self.relu(self.conv1_1(mean_sub_input)))) | |
| main_out_pool1 = self.pool1(main_out_block1) | |
| main_out_block2 = self.relu(self.conv2_2(self.relu(self.conv2_1(main_out_pool1)))) | |
| main_out_pool2 = self.pool2(main_out_block2) | |
| main_out_block3 = self.relu(self.conv3_3(self.relu(self.conv3_2(self.relu(self.conv3_1(main_out_pool2)))))) | |
| main_out_pool3 = self.pool3(main_out_block3) | |
| main_out_block4 = self.relu(self.conv4_3(self.relu(self.conv4_2(self.relu(self.conv4_1(main_out_pool3)))))) | |
| main_out_pool4 = self.pool3(main_out_block4) | |
| main_out_block5 = self.relu(self.conv_before_transpose_1( | |
| self.relu(self.conv5_3(self.relu(self.conv5_2(self.relu(self.conv5_1(main_out_pool4)))))))) | |
| main_out_rest = self.convA_5(self.relu( | |
| self.convA_4(self.relu(self.convA_3(self.relu(self.convA_2(self.relu(self.convA_1(main_out_block5))))))))) | |
| if self.name == "scale_1": | |
| return main_out_rest | |
| ################## Stage 2 ############################ | |
| sub1_out_conv1 = self.relu(self.conv_mid_4(self.relu( | |
| self.conv_middle_3(self.relu(self.conv_middle_2(self.relu(self.conv_middle_1(main_out_pool3)))))))) | |
| sub1_transpose = self.relu(self.transpose_1(main_out_block5)) | |
| sub1_after_transpose_1 = self.relu(self.conv_after_transpose_1_1(sub1_transpose)) | |
| sub1_concat = torch.cat((sub1_out_conv1, sub1_after_transpose_1), dim=1) | |
| sub1_out_rest = self.convB_5(self.relu( | |
| self.convB_4(self.relu(self.convB_3(self.relu(self.convB_2(self.relu(self.convB_1(sub1_concat))))))))) | |
| if self.name == "scale_2": | |
| return main_out_rest, sub1_out_rest | |
| ################# Stage 3 ############################ | |
| sub2_out_conv1 = self.relu(self.conv_lowest_4(self.relu( | |
| self.conv_lowest_3(self.relu(self.conv_lowest_2(self.relu(self.conv_lowest_1(main_out_pool2)))))))) | |
| sub2_transpose = self.relu(self.transpose_2(sub1_out_conv1)) | |
| sub2_after_transpose_1 = self.relu(self.conv_after_transpose_2_1(sub2_transpose)) | |
| sub3_transpose = self.relu(self.transpose_3(main_out_block5)) | |
| sub3_after_transpose_1 = self.relu(self.conv_after_transpose_3_1(sub3_transpose)) | |
| sub2_concat = torch.cat((sub2_out_conv1, sub2_after_transpose_1, sub3_after_transpose_1), dim=1) | |
| sub2_out_rest = self.convC_5(self.relu( | |
| self.convC_4(self.relu(self.convC_3(self.relu(self.convC_2(self.relu(self.convC_1(sub2_concat))))))))) | |
| if self.name == "scale_3": | |
| return main_out_rest, sub1_out_rest, sub2_out_rest | |
| ################# Stage 4 ############################ | |
| sub4_out_conv1 = self.relu( | |
| self.conv_scale1_3(self.relu(self.conv_scale1_2(self.relu(self.conv_scale1_1(main_out_pool1)))))) | |
| # TDF 1 | |
| tdf_4_1_a = self.relu(self.transpose_4_1_a(main_out_block5)) | |
| tdf_4_1_b = self.relu(self.transpose_4_1_b(tdf_4_1_a)) | |
| after_tdf_4_1 = self.relu(self.conv_after_transpose_4_1(tdf_4_1_b)) | |
| # TDF 2 | |
| tdf_4_2 = self.relu(self.transpose_4_2(sub1_out_conv1)) | |
| after_tdf_4_2 = self.relu(self.conv_after_transpose_4_2(tdf_4_2)) | |
| # TDF 3 | |
| tdf_4_3 = self.relu(self.transpose_4_3(sub2_out_conv1)) | |
| after_tdf_4_3 = self.relu(self.conv_after_transpose_4_3(tdf_4_3)) | |
| sub4_concat = torch.cat((sub4_out_conv1, after_tdf_4_1, after_tdf_4_2, after_tdf_4_3), dim=1) | |
| sub4_out_rest = self.convD_5(self.relu( | |
| self.convD_4(self.relu(self.convD_3(self.relu(self.convD_2(self.relu(self.convD_1(sub4_concat))))))))) | |
| logging.info("Forward Finished") | |
| if self.name == "scale_4": | |
| return main_out_rest, sub1_out_rest, sub2_out_rest, sub4_out_rest | |
| def predict_single_image(self, image, emoji, nms_thresh=0.25, thickness=2, multi_colours=True): | |
| if image.shape[0] % 16 or image.shape[1] % 16: | |
| image = cv2.resize(image, (image.shape[1]//16*16, image.shape[0]//16*16)) | |
| img_tensor = torch.from_numpy(image.transpose((2, 0, 1)).astype(np.float32)).unsqueeze(0) | |
| with torch.no_grad(): | |
| out = self.forward(img_tensor.cuda()) | |
| # out = self.forward(img_tensor) | |
| out = get_upsample_output(out, self.output_downscale) | |
| pred_dot_map, pred_box_map = get_box_and_dot_maps(out, nms_thresh, self.BOXES) | |
| img_out = get_boxed_img(image, emoji, pred_box_map, pred_box_map, pred_dot_map, self.output_downscale, | |
| self.BOXES, self.BOX_SIZE_BINS, thickness=thickness, multi_colours=multi_colours) | |
| return pred_dot_map, pred_box_map, img_out |