| 44 | self.num_steps = self.config.Global.batch_max_length + 1 |
| 45 | |
| 46 | def preprocess(self, image): |
| 47 | self.transform = transforms.ToTensor() |
| 48 | |
| 49 | if self.keep_ratio_with_pad: |
| 50 | w, h = image.size |
| 51 | ratio = w / float(h) |
| 52 | if math.ceil(ratio * self.imgH) > self.imgW: |
| 53 | resized_image = image.resize((self.imgW, self.imgH), Image.BICUBIC) |
| 54 | resized_image = self.transform(resized_image) |
| 55 | imgP = resized_image.sub(0.5).div(0.5) |
| 56 | else: |
| 57 | resized_W = math.ceil(ratio * self.imgH) |
| 58 | resized_image = image.resize((resized_W, self.imgH), Image.BICUBIC) |
| 59 | resized_image = self.transform(resized_image) |
| 60 | resized_image = resized_image.sub(0.5).div(0.5) |
| 61 | |
| 62 | c, h, w = resized_image.size() |
| 63 | imgP = torch.FloatTensor(*(self.channel, self.imgH, self.imgW)).fill_(0) |
| 64 | imgP[:, :, :w] = resized_image |
| 65 | imgP[:, :, w:] = resized_image[:, :, w - 1].unsqueeze(2).expand(c, h, self.imgW - w) |
| 66 | else: |
| 67 | resized_image = image.resize((self.imgW, self.imgH), Image.BICUBIC) |
| 68 | resized_image = self.transform(resized_image) |
| 69 | imgP = resized_image.sub(0.5).div(0.5) |
| 70 | |
| 71 | imgP = imgP.unsqueeze(0) |
| 72 | return imgP |
| 73 | |
| 74 | def predict(self, image_tensor): |
| 75 | self.model.eval() |