def __init__(self,
class_id=0,
track_id=0,
polygon=Polygon(),
bbox=Bbox(),
keypoints=Keypoints()):
self.class_id = class_id
self.track_id = track_id
self.polygon = polygon
self.bbox = bbox
self.keypoints = keypoints
def detect_single_image(self, image_path, crop_area=None):
img = Image.open(image_path).convert('RGB')
W, H = img.size
# Cropping
if crop_area is None:
crop_area = (0, 0, W, H)
x_crop, y_crop, w_crop, h_crop = crop_area
margin = 0.03 * max(W, H) # 3% margin
detections = []
# Number of repeated cropping areas to span the entire image
n_left = math.ceil(x_crop / w_crop)
n_right = math.ceil((W - (x_crop + w_crop)) / w_crop)
n_top = math.ceil(y_crop / h_crop)
n_bottom = math.ceil((H - (y_crop + h_crop)) / h_crop)
for i in range(-n_top, 1 + n_bottom):
for j in range(-n_left, 1 + n_right):
crop_area = (x_crop + j * w_crop - margin,
y_crop + i * h_crop - margin,
x_crop + j * w_crop + w_crop + margin,
y_crop + i * h_crop + h_crop + margin)
print(i, j, crop_area)
img_crop = img.crop(crop_area)
sized = letterbox_image(img_crop, self.model.width,
self.model.height)
start = time.time()
with torch.no_grad():
boxes = do_detect(self.model, sized, 0.5, 0.4,
self.use_cuda)
self.correct_yolo_boxes(boxes, W, H, self.model.width,
self.model.height, crop_area)
finish = time.time()
print('%s: Predicted in %f seconds.' % (image_path,
(finish - start)))
detections.extend([
Detection(class_id=self.yolo_to_coco_class(b[6].item()),
track_id=i,
bbox=Bbox((b[0] - b[2] / 2).item() * W,
(b[1] - b[3] / 2).item() * H,
b[2].item() * W, b[3].item() * H))
for i, b in enumerate(boxes)
])
return detections
def keypoints_to_bbox(kps):
x, y, v = kps[0::3], kps[1::3], kps[2::3]
x, y = x[v > 0], y[v > 0]
x1, y1, x2, y2 = np.min(x), np.min(y), np.max(x), np.max(y)
return Bbox(x1, y1, x2 - x1, y2 - y1)
def receive_bbox(self):
data = self.request.recv(1024)
bbox = pickle.loads(data)
return Bbox(*bbox)
def from_df(row):
bbox = Bbox(row.x, row.y, row.w, row.h)
return Detection(row.class_id, row.track_id,
Polygon.from_str(row.polygon), bbox,
Keypoints.from_str(row.kp))
def from_json(data):
return Detection(data["class_id"], data["track_id"],
Polygon(data["polygon"]), Bbox(*data["bbox"]),
Keypoints(data["keypoints"]))