| import os | |
| import cv2 | |
| import torch | |
| import argparse | |
| from tqdm import tqdm | |
| from huggingface_hub import PyTorchModelHubMixin | |
| from ddcolor_model import DDColor | |
| from infer import ImageColorizationPipeline | |
| class DDColorHF(DDColor, PyTorchModelHubMixin): | |
| def __init__(self, config): | |
| super().__init__(**config) | |
| class ImageColorizationPipelineHF(ImageColorizationPipeline): | |
| def __init__(self, model, input_size): | |
| self.input_size = input_size | |
| if torch.cuda.is_available(): | |
| self.device = torch.device("cuda") | |
| else: | |
| self.device = torch.device("cpu") | |
| self.model = model.to(self.device) | |
| self.model.eval() | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--model_name", type=str, default="ddcolor_modelscope") | |
| parser.add_argument( | |
| "--input", | |
| type=str, | |
| default="figure/", | |
| help="input test image folder or video path", | |
| ) | |
| parser.add_argument( | |
| "--output", type=str, default="results", help="output folder or video path" | |
| ) | |
| parser.add_argument( | |
| "--input_size", type=int, default=512, help="input size for model" | |
| ) | |
| args = parser.parse_args() | |
| if not os.path.exists(args.model_name): | |
| model_name = f"piddnad/{args.model_name}" | |
| else: | |
| model_name = args.model_name | |
| ddcolor_model = DDColorHF.from_pretrained(model_name) | |
| print(f"Output path: {args.output}") | |
| os.makedirs(args.output, exist_ok=True) | |
| img_list = os.listdir(args.input) | |
| assert len(img_list) > 0 | |
| colorizer = ImageColorizationPipelineHF( | |
| model=ddcolor_model, input_size=args.input_size | |
| ) | |
| for name in tqdm(img_list): | |
| img = cv2.imread(os.path.join(args.input, name)) | |
| image_out = colorizer.process(img) | |
| cv2.imwrite(os.path.join(args.output, name), image_out) | |
| if __name__ == "__main__": | |
| main() | |