()
| 54 | cv2.imwrite(osp.join(args.output, f'{imgname}_{args.model}.png'), output) |
| 55 | |
| 56 | def main(): |
| 57 | parser = argparse.ArgumentParser() |
| 58 | parser.add_argument('--input', type=str, required=True,help='input test image folder') |
| 59 | parser.add_argument('--output', type=str, default='outputs/', help='output folder') |
| 60 | parser.add_argument('--model',type=str,default='ram_promptir', help='model type') |
| 61 | # TODO: it now only supports sr, need to adapt to dn and jpeg_car |
| 62 | parser.add_argument('--model_path',type=str,default='pretrained_model/ram_promptir_finetune.pth') |
| 63 | args = parser.parse_args() |
| 64 | |
| 65 | os.makedirs(args.output, exist_ok=True) |
| 66 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| 67 | # set up model |
| 68 | model = define_model(args) |
| 69 | model.eval() |
| 70 | model = model.to(device) |
| 71 | |
| 72 | if osp.isdir(args.input): |
| 73 | os.makedirs(args.output, exist_ok=True) |
| 74 | for idx, path in enumerate(sorted(glob.glob(osp.join(args.input, '*')))): |
| 75 | if path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif')): |
| 76 | process_image(path, model, device, args) |
| 77 | elif osp.isfile(args.input) and args.input.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif')): |
| 78 | os.makedirs(args.output, exist_ok=True) |
| 79 | process_image(args.input, model, device, args) |
| 80 | else: |
| 81 | print('invalid img format') |
| 82 | |
| 83 | |
| 84 | if __name__ == "__main__": |
no test coverage detected