| 19 | pass |
| 20 | |
| 21 | def test_model(self): |
| 22 | # Images |
| 23 | imgs = [ |
| 24 | "backend/test_images/bus.jpg", |
| 25 | "backend/test_images/zidane.jpg", |
| 26 | "backend/test_images/bear.jpg", |
| 27 | "backend/test_images/dragon.jpg", |
| 28 | ] |
| 29 | img = Image.open(imgs[3]) |
| 30 | |
| 31 | # Inference |
| 32 | results = model(img) |
| 33 | self.assertIsNotNone(results) |
| 34 | |
| 35 | # Results |
| 36 | # results.print() |
| 37 | # results.save() |
| 38 | # print(results.xyxy[0]) # print img1 predictions (pixels) |
| 39 | # results.show() |
| 40 | print(results.names) |
| 41 | |
| 42 | # https://github.com/ultralytics/yolov5/blob/master/models/common.py |
| 43 | names = results.names |
| 44 | detected = [] |
| 45 | if results.pred is not None: |
| 46 | pred = results.pred[0] |
| 47 | if pred is not None: |
| 48 | for c in pred[:, -1].unique(): |
| 49 | n = (pred[:, -1] == c).sum() |
| 50 | detected.append(f"{n} {names[int(c)]}{'s' * (n > 1)}") |
| 51 | |
| 52 | print(f"Detected: {detected}") |
| 53 | |
| 54 | processed_imgs = results.render() |
| 55 | processed_img = Image.fromarray(processed_imgs[0]).convert("RGB") |
| 56 | processed_img.save("test.jpg") |