| 156 | return boxes |
| 157 | |
| 158 | def predict(self, image, imgsz=1024, **kwargs): |
| 159 | # Preprocess input image |
| 160 | orig_h, orig_w = image.shape[:2] |
| 161 | pix = self.resize_and_pad_image(image, new_shape=imgsz) |
| 162 | pix = np.transpose(pix, (2, 0, 1)) # CHW |
| 163 | pix = np.expand_dims(pix, axis=0) # BCHW |
| 164 | pix = pix.astype(np.float32) / 255.0 # Normalize to [0, 1] |
| 165 | new_h, new_w = pix.shape[2:] |
| 166 | |
| 167 | # Run inference |
| 168 | preds = self.model.run(None, {"images": pix})[0] |
| 169 | |
| 170 | # Postprocess predictions |
| 171 | preds = preds[preds[..., 4] > 0.25] |
| 172 | preds[..., :4] = self.scale_boxes( |
| 173 | (new_h, new_w), preds[..., :4], (orig_h, orig_w) |
| 174 | ) |
| 175 | return [YoloResult(boxes=preds, names=self._names)] |
| 176 | |
| 177 | |
| 178 | class ModelInstance: |