Update Self-Correction-Human-Parsing/datasets/simple_extractor_dataset.py

Self-Correction-Human-Parsing/datasets/simple_extractor_dataset.py

@@ -1,36 +1,18 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   [email protected]
-@File    :   dataset.py
-@Time    :   8/30/19 9:12 PM
-@Desc    :   Dataset Definition
-@License :   This source code is licensed under the license found in the
-             LICENSE file in the root directory of this source tree.
-"""
-
-import os
 import cv2
 import numpy as np
-
 from torch.utils import data
 from utils.transforms import get_affine_transform
 
-
-class SimpleFolderDataset(data.Dataset):
-    def __init__(self, root, input_size=[512, 512], transform=None):
-        self.root = root
+class SimpleFileDataset(data.Dataset):
+    def __init__(self, img_path, input_size=[512, 512], transform=None):
+        self.img_path = img_path  # A single file path
         self.input_size = input_size
         self.transform = transform
         self.aspect_ratio = input_size[1] * 1.0 / input_size[0]
         self.input_size = np.asarray(input_size)
 
-        self.file_list = os.listdir(self.root)
-
     def __len__(self):
-        return len(self.file_list)
+        return 1  # Only one image, so the length is 1
 
     def _box2cs(self, box):
         x, y, w, h = box[:4]
@@ -48,16 +30,14 @@ class SimpleFolderDataset(data.Dataset):
         return center, scale
 
     def __getitem__(self, index):
-        img_name = self.file_list[index]
-        img_path = os.path.join(self.root, img_name)
-        img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+        img = cv2.imread(self.img_path, cv2.IMREAD_COLOR)
         h, w, _ = img.shape
 
         # Get person center and scale
         person_center, s = self._box2cs([0, 0, w - 1, h - 1])
         r = 0
         trans = get_affine_transform(person_center, s, r, self.input_size)
-        input = cv2.warpAffine(
+        input_img = cv2.warpAffine(
             img,
             trans,
             (int(self.input_size[1]), int(self.input_size[0])),
@@ -65,9 +45,9 @@ class SimpleFolderDataset(data.Dataset):
             borderMode=cv2.BORDER_CONSTANT,
             borderValue=(0, 0, 0))
 
-        input = self.transform(input)
+        input_img = self.transform(input_img)
         meta = {
-            'name': img_name,
+            'name': self.img_path,
             'center': person_center,
             'height': h,
             'width': w,
@@ -75,4 +55,4 @@ class SimpleFolderDataset(data.Dataset):
             'rotation': r
         }
 
-        return input, meta
+        return input_img, meta