Update Self-Correction-Human-Parsing/simple_extractor.py

Self-Correction-Human-Parsing/simple_extractor.py

@@ -10,7 +10,7 @@ import torchvision.transforms as transforms
 
 import networks
 from utils.transforms import transform_logits
-from datasets.simple_extractor_dataset import SimpleFolderDataset
+from datasets.simple_extractor_dataset import SimpleFileDataset  # Modify dataset class to use SimpleFileDataset
 
 dataset_settings = {
     'atr': {
@@ -27,7 +27,7 @@ def get_arguments():
     parser.add_argument("--dataset", type=str, default='atr', choices=['atr'])
     parser.add_argument("--model-restore", type=str, default='', help="Path to pretrained model.")
     parser.add_argument("--gpu", type=str, default='0', help="GPU device.")
-    parser.add_argument("--input-dir", type=str, default='', help="Path of input image folder.")
+    parser.add_argument("--input-path", type=str, default='', help="Path to a single input image.")
     parser.add_argument("--output-dir", type=str, default='', help="Path of output image folder.")
 
     return parser.parse_args()
@@ -57,8 +57,9 @@ def main():
         transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229])
     ])
     
-    dataset = SimpleFolderDataset(root=args.input_dir, input_size=input_size, transform=transform)
-    dataloader = DataLoader(dataset)
+    # Use the SimpleFileDataset class instead of SimpleFolderDataset
+    dataset = SimpleFileDataset(img_path=args.input_path, input_size=input_size, transform=transform)
+    dataloader = DataLoader(dataset, batch_size=1)  # Only one image, so batch_size=1
 
     if not os.path.exists(args.output_dir):
         os.makedirs(args.output_dir)