示例#1
0
    def __getitem__(self, idx):

        image_path = os.path.join(self.data_root,
                                  self.landmarks_frame.iloc[idx, 0])
        scale = self.landmarks_frame.iloc[idx, 1]

        center_w = self.landmarks_frame.iloc[idx, 2]
        center_h = self.landmarks_frame.iloc[idx, 3]
        center = torch.Tensor([center_w, center_h])

        pts = self.landmarks_frame.iloc[idx, 4:].values
        pts = pts.astype('float').reshape(-1, 2)

        scale *= 1.25
        nparts = pts.shape[0]
        img = np.array(Image.open(image_path).convert('RGB'), dtype=np.float32)

        r = 0
        if self.is_train:
            scale = scale * (random.uniform(1 - self.scale_factor,
                                            1 + self.scale_factor))
            r = random.uniform(-self.rot_factor, self.rot_factor) \
                if random.random() <= 0.6 else 0
            if random.random() <= 0.5 and self.flip:
                img = np.fliplr(img)
                pts = fliplr_joints(pts, width=img.shape[1], dataset='WFLW')
                center[0] = img.shape[1] - center[0]

        img = crop(img, center, scale, self.input_size, rot=r)

        target = np.zeros((nparts, self.output_size[0], self.output_size[1]))
        tpts = pts.copy()

        for i in range(nparts):
            if tpts[i, 1] > 0:
                tpts[i, 0:2] = transform_pixel(tpts[i, 0:2] + 1,
                                               center,
                                               scale,
                                               self.output_size,
                                               rot=r)
                target[i] = generate_target(target[i],
                                            tpts[i] - 1,
                                            self.sigma,
                                            label_type=self.label_type)
        img = img.astype(np.float32)
        img = (img / 255.0 - self.mean) / self.std
        img = img.transpose([2, 0, 1])
        target = torch.Tensor(target)
        tpts = torch.Tensor(tpts)
        center = torch.Tensor(center)

        meta = {
            'index': idx,
            'center': center,
            'scale': scale,
            'pts': torch.Tensor(pts),
            'tpts': tpts
        }

        return img, target, meta
示例#2
0
    def __getitem__(self, idx):

        img = self.images[idx][0]

        if len(img.shape) == 2:
            img = img.reshape(img.shape[0], img.shape[1], 1)
            img = np.repeat(img, 3, axis=2)

        pts = self.pts[idx][0:58].reshape(2, -1).transpose()

        xmin = np.min(pts[:, 0])
        xmax = np.max(pts[:, 0])
        ymin = np.min(pts[:, 1])
        ymax = np.max(pts[:, 1])

        center_w = (math.floor(xmin) + math.ceil(xmax)) / 2.0
        center_h = (math.floor(ymin) + math.ceil(ymax)) / 2.0

        scale = max(math.ceil(xmax) - math.floor(xmin), math.ceil(ymax) - math.floor(ymin)) / 200.0
        center = torch.Tensor([center_w, center_h])

        scale *= 1.25
        nparts = pts.shape[0]

        r = 0
        if self.is_train:
            scale = scale * (random.uniform(1 - self.scale_factor,
                                            1 + self.scale_factor))
            r = random.uniform(-self.rot_factor, self.rot_factor) \
                if random.random() <= 0.6 else 0

            if random.random() <= 0.5 and self.flip:
                img = np.fliplr(img)
                pts = fliplr_joints(pts, width=img.shape[1], dataset='COFW')
                center[0] = img.shape[1] - center[0]

        img = crop(img, center, scale, self.input_size, rot=r)

        target = np.zeros((nparts, self.output_size[0], self.output_size[1]))
        tpts = pts.copy()

        for i in range(nparts):
            if tpts[i, 1] > 0:
                tpts[i, 0:2] = transform_pixel(tpts[i, 0:2]+1, center,
                                               scale, self.output_size, rot=r)
                target[i] = generate_target(target[i], tpts[i]-1, self.sigma,
                                            label_type=self.label_type)
        img = img.astype(np.float32)
        img = (img/255 - self.mean) / self.std
        img = img.transpose([2, 0, 1])
        target = torch.Tensor(target)
        tpts = torch.Tensor(tpts)
        center = torch.Tensor(center)

        meta = {'index': idx, 'center': center, 'scale': scale,
                'pts': torch.Tensor(pts), 'tpts': tpts}

        return img, target, meta
def get_training_image(img_path,
                       bbox=None,
                       inp_res=256,
                       mean=(0.6419, 0.6292, 0.5994),
                       std=(0.2311, 0.2304, 0.2379)):
    img = load_image(img_path)

    if bbox is not None:
        x0, y0, x1, y1 = bbox[0][0], bbox[0][1], bbox[1][0], bbox[1][1]
        c = np.array([(x0 + x1), (y0 + y1)]) / 2  # center
        s = np.sqrt((y1 - y0) * (x1 - x0)) / 60.0  # scale

    else:
        c = np.array([img.shape[2] / 2, img.shape[1] / 2])
        s = 5.0  # THIS HAS TO BE FIXED !!!

    r = 0  # rotation

    inp = crop(img, c, s, [inp_res, inp_res], rot=r)
    inp = color_normalize(inp, mean, std)

    meta = {'center': c, 'scale': s}

    return inp, meta
示例#4
0
    def __getitem__(self, idx):
        if self.landmarks_frame.iloc[idx, 0].find('/ibug/image_092_01.jpg'):
            self.landmarks_frame.iloc[idx, 0] = self.landmarks_frame.iloc[
                idx, 0].replace('image_092_01.jpg', 'image_092 _01.jpg')

        image_path = os.path.join(self.data_root,
                                  self.landmarks_frame.iloc[idx, 0])

        scale = self.landmarks_frame.iloc[idx, 1]

        center_w = self.landmarks_frame.iloc[idx, 2]
        center_h = self.landmarks_frame.iloc[idx, 3]
        center = torch.Tensor([center_w, center_h])

        pts = self.landmarks_frame.iloc[idx, 4:].values
        pts = pts.astype('float').reshape(-1, 2)

        scale *= 1.25
        nparts = pts.shape[0]
        img = np.array(Image.open(image_path).convert('RGB'), dtype=np.float32)

        r = 0
        img = crop(img, center, scale, self.input_size, rot=r)

        trans_imgs = []
        trans_factors = []
        for _ in range(self.n_trans):
            x_factor = random.uniform(-self.trans_factor, self.trans_factor)
            y_factor = random.uniform(-self.trans_factor, self.trans_factor)

            M = np.float32([[1, 0, self.input_size[0] * x_factor],
                            [0, 1, self.input_size[1] * y_factor]])
            trans_img = cv2.warpAffine(img, M, self.input_size)

            trans_img = trans_img.astype(np.float32)
            trans_img = (trans_img / 255.0 - self.mean) / self.std
            trans_img = trans_img.transpose([2, 0, 1])

            trans_imgs.append(trans_img)
            trans_factors.append([x_factor, y_factor])

        img = img.astype(np.float32)
        img = (img / 255.0 - self.mean) / self.std
        img = img.transpose([2, 0, 1])

        targets = []
        offsets = []
        for output_size in self.output_size:
            target = np.zeros((nparts, output_size[0], output_size[1]))
            tpts = pts.copy()

            for i in range(nparts):
                if tpts[i, 1] > 0:
                    tpts[i, 0:2] = transform_pixel(tpts[i, 0:2],
                                                   center,
                                                   scale,
                                                   output_size,
                                                   rot=r)
                    x, y = tpts[i].astype(int)

                    if 0 <= y < output_size[0] and 0 <= x < output_size[1]:
                        target[i, y, x] = 1

            xx_channel, yy_channel = self._generate_offset(
                nparts, tpts, output_size, self.offset_mode, self.offset_dim)
            offset = [torch.Tensor(xx_channel), torch.Tensor(yy_channel)]

            targets.append(torch.Tensor(target))
            offsets.append(offset)

        trans_imgs = torch.Tensor(trans_imgs)
        trans_factors = torch.Tensor(trans_factors)
        tpts = torch.Tensor(tpts)
        center = torch.Tensor(center)

        meta = {
            'index': idx,
            'center': center,
            'scale': scale,
            'pts': torch.Tensor(pts),
            'tpts': tpts
        }

        return img, trans_imgs, trans_factors, targets, offsets, meta
示例#5
0
    def __getitem__(self, index):
        sf = self.scale_factor
        rf = self.rot_factor
        if self.is_train:
            a = self.anno[self.train_list[index]]
        else:
            a = self.anno[self.valid_list[index]]

        img_path = os.path.join(self.img_folder, a['img_paths'])
        pts = torch.Tensor(a['joint_self'])
        # pts[:, 0:2] -= 1  # Convert pts to zero based

        # c = torch.Tensor(a['objpos']) - 1
        c = torch.Tensor(a['objpos'])
        s = a['scale_provided']

        # Adjust center/scale slightly to avoid cropping limbs
        if c[0] != -1:
            c[1] = c[1] + 15 * s
            s = s * 1.25

        # For single-person pose estimation with a centered/scaled figure
        nparts = pts.size(0)
        img = load_image(img_path)  # CxHxW

        r = 0
        if self.is_train:
            s = s * torch.randn(1).mul_(sf).add_(1).clamp(1 - sf, 1 + sf)[0]
            r = torch.randn(1).mul_(rf).clamp(
                -2 * rf, 2 * rf)[0] if random.random() <= 0.6 else 0

            # Flip
            if random.random() <= 0.5:
                img = fliplr(img)
                pts = shufflelr(pts, img.size(2), self.DATA_INFO.hflip_indices)
                c[0] = img.size(2) - c[0]

            # Color
            img[0, :, :].mul_(random.uniform(0.8, 1.2)).clamp_(0, 1)
            img[1, :, :].mul_(random.uniform(0.8, 1.2)).clamp_(0, 1)
            img[2, :, :].mul_(random.uniform(0.8, 1.2)).clamp_(0, 1)

        # Prepare image and groundtruth map
        inp = crop(img, c, s, self.inp_res, rot=r)
        inp = color_normalize(inp, self.DATA_INFO.rgb_mean,
                              self.DATA_INFO.rgb_stddev)

        # Generate ground truth
        tpts = pts.clone()
        target = torch.zeros(nparts, *self.out_res)
        target_weight = tpts[:, 2].clone().view(nparts, 1)

        for i in range(nparts):
            # if tpts[i, 2] > 0: # This is evil!!
            if tpts[i, 1] > 0:
                tpts[i, 0:2] = to_torch(
                    transform(tpts[i, 0:2] + 1, c, s, self.out_res, rot=r))
                target[i], vis = draw_labelmap(target[i],
                                               tpts[i] - 1,
                                               self.sigma,
                                               type=self.label_type)
                target_weight[i, 0] *= vis

        # Meta info
        if not isinstance(s, torch.Tensor):
            s = torch.Tensor(s)

        meta = {
            'index': index,
            'center': c,
            'scale': s,
            'pts': pts,
            'tpts': tpts,
            'target_weight': target_weight
        }

        return inp, target, meta
示例#6
0
    def __getitem__(self, idx):
        if self.landmarks_frame.iloc[idx, 0].find('/ibug/image_092_01.jpg'):
            self.landmarks_frame.iloc[idx, 0] = self.landmarks_frame.iloc[
                idx, 0].replace('image_092_01.jpg', 'image_092 _01.jpg')

        image_path = os.path.join(self.data_root,
                                  self.landmarks_frame.iloc[idx, 0])
        scale = self.landmarks_frame.iloc[idx, 1]

        center_w = self.landmarks_frame.iloc[idx, 2]
        center_h = self.landmarks_frame.iloc[idx, 3]
        center = torch.Tensor([center_w, center_h])

        pts = self.landmarks_frame.iloc[idx, 4:].values
        pts = pts.astype('float').reshape(-1, 2)

        scale *= 1.3
        nparts = pts.shape[0]
        img = np.array(Image.open(image_path).convert('RGB'), dtype=np.float32)

        r = 0
        t_x, t_y = (1, 1)
        if self.is_train:
            scale = scale * (random.uniform(1 - self.scale_factor,
                                            1 + self.scale_factor))
            r = random.uniform(-self.rot_factor, self.rot_factor) \
                if random.random() <= 0.6 else 0
            t_x = random.uniform(1-self.trans_factor, 1+self.trans_factor) \
                if random.random() <= 0.5 else 1
            t_y = random.uniform(1-self.trans_factor, 1+self.trans_factor) \
                if random.random() <= 0.5 else 1
            if random.random() <= 0.5 and self.flip:
                img = np.fliplr(img)
                pts = fliplr_joints(pts, width=img.shape[1], dataset='300W')
                center[0] = img.shape[1] - center[0]

        img = crop(img,
                   center,
                   scale,
                   self.input_size,
                   translation=(t_x, t_y),
                   rot=r)

        if self.is_train:
            if random.random() <= 0.3 and self.gaussian_blur:
                radius = random.choice([1, 3, 5])
                img = cv2.GaussianBlur(img, (radius, radius), sigmaX=1.0)
            if random.random() <= 1.0 and self.occlusion:
                img = add_occlusion(img, max_size=102)

        targets = []
        offsets = []
        for output_size in self.output_size:
            target = np.zeros((nparts, output_size[0], output_size[1]))
            tpts = pts.copy()

            for i in range(nparts):
                if tpts[i, 1] > 0:
                    tpts[i, 0:2] = transform_pixel(tpts[i, 0:2],
                                                   center,
                                                   scale,
                                                   output_size,
                                                   translation=(t_x, t_y),
                                                   rot=r)
                    x, y = tpts[i].astype(int)

                    if 0 <= y < output_size[0] and 0 <= x < output_size[1]:
                        target[i, y, x] = 1

            xx_channel, yy_channel = self._generate_offset(
                nparts, tpts, output_size, self.offset_mode, self.offset_dim)
            offset = [torch.Tensor(xx_channel), torch.Tensor(yy_channel)]

            targets.append(torch.Tensor(target))
            offsets.append(offset)

        img = img.astype(np.float32)
        img = (img / 255.0 - self.mean) / self.std
        img = img.transpose([2, 0, 1])
        tpts = torch.Tensor(tpts)
        center = torch.Tensor(center)

        meta = {
            'index': idx,
            'center': center,
            'scale': scale,
            'pts': torch.Tensor(pts),
            'tpts': tpts
        }

        return img, targets, offsets, meta