| 45 | |
| 46 | |
| 47 | def postprcess(out): |
| 48 | numClasses = 20 |
| 49 | anchors = [1.08, 1.19, 3.42, 4.41, 6.63, 11.38, 9.42, 5.11, 16.62, 10.52] |
| 50 | |
| 51 | def sigmoid(x, derivative=False): |
| 52 | return x * (1 - x) if derivative else 1 / (1 + np.exp(-x)) |
| 53 | |
| 54 | def softmax(x): |
| 55 | scoreMatExp = np.exp(np.asarray(x)) |
| 56 | return scoreMatExp / scoreMatExp.sum(0) |
| 57 | |
| 58 | clut = [(0, 0, 0), (255, 0, 0), (255, 0, 255), (0, 0, 255), (0, 255, 0), |
| 59 | (0, 255, 128), (128, 255, 0), (128, 128, 0), (0, 128, 255), |
| 60 | (128, 0, 128), (255, 0, 128), (128, 0, 255), (255, 128, 128), |
| 61 | (128, 255, 128), (255, 255, 0), (255, 128, 128), (128, 128, 255), |
| 62 | (255, 128, 128), (128, 255, 128)] |
| 63 | label = [ |
| 64 | "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", |
| 65 | "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", |
| 66 | "pottedplant", "sheep", "sofa", "train", "tvmonitor" |
| 67 | ] |
| 68 | |
| 69 | img = get_image() |
| 70 | draw = ImageDraw.Draw(img) |
| 71 | |
| 72 | for cy in range(13): |
| 73 | for cx in range(13): |
| 74 | for b in range(5): |
| 75 | channel = b * (numClasses + 5) |
| 76 | tx = out[channel][cy][cx] |
| 77 | ty = out[channel + 1][cy][cx] |
| 78 | tw = out[channel + 2][cy][cx] |
| 79 | th = out[channel + 3][cy][cx] |
| 80 | tc = out[channel + 4][cy][cx] |
| 81 | x = (float(cx) + sigmoid(tx)) * 32 |
| 82 | y = (float(cy) + sigmoid(ty)) * 32 |
| 83 | |
| 84 | w = np.exp(tw) * 32 * anchors[2 * b] |
| 85 | h = np.exp(th) * 32 * anchors[2 * b + 1] |
| 86 | |
| 87 | confidence = sigmoid(tc) |
| 88 | |
| 89 | classes = np.zeros(numClasses) |
| 90 | for c in range(0, numClasses): |
| 91 | classes[c] = out[channel + 5 + c][cy][cx] |
| 92 | |
| 93 | classes = softmax(classes) |
| 94 | detectedClass = classes.argmax() |
| 95 | if 0.5 < classes[detectedClass] * confidence: |
| 96 | color = clut[detectedClass] |
| 97 | x = x - w / 2 |
| 98 | y = y - h / 2 |
| 99 | draw.line((x, y, x + w, y), fill=color) |
| 100 | draw.line((x, y, x, y + h), fill=color) |
| 101 | draw.line((x + w, y, x + w, y + h), fill=color) |
| 102 | draw.line((x, y + h, x + w, y + h), fill=color) |
| 103 | draw.text((x, y), label[detectedClass], fill=color) |
| 104 | logging.info("bounding box: (%.2f, %.2f, %.2f, %.2f)" % |