(input_path=None, output_path=None, aligned=False, bg_sr=False, scale_factor=2, save_text=False, device=None)
| 13 | |
| 14 | |
| 15 | def inference(input_path=None, output_path=None, aligned=False, bg_sr=False, scale_factor=2, save_text=False, device=None): |
| 16 | |
| 17 | if device == None or device == 'gpu': |
| 18 | use_cuda = torch.cuda.is_available() |
| 19 | if device == 'cpu': |
| 20 | use_cuda = False |
| 21 | device = torch.device('cuda' if use_cuda else 'cpu') |
| 22 | |
| 23 | if input_path is None: |
| 24 | exit('input image path is none. Please see our document') |
| 25 | |
| 26 | if output_path is None: |
| 27 | TIMESTAMP = time.strftime("%m-%d_%H-%M", time.localtime()) |
| 28 | if input_path[-1] == '/' or input_path[-1] == '\\': |
| 29 | input_path = input_path[:-1] |
| 30 | output_path = osp.join(input_path+'_'+TIMESTAMP+'_MARCONetPlus') |
| 31 | os.makedirs(output_path, exist_ok=True) |
| 32 | |
| 33 | # use bsrgan to restore the background of the whole image |
| 34 | if bg_sr: |
| 35 | ##Define BG restoration model |
| 36 | BGModel = BSRGAN(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=2) # define network |
| 37 | model_old = torch.load('./checkpoints/bsrgan_bg.pth') |
| 38 | state_dict = BGModel.state_dict() |
| 39 | for ((key, param),(key2, param2)) in zip(model_old.items(), state_dict.items()): |
| 40 | state_dict[key2] = param |
| 41 | BGModel.load_state_dict(state_dict, strict=True) |
| 42 | BGModel.eval() |
| 43 | for k, v in BGModel.named_parameters(): |
| 44 | v.requires_grad = False |
| 45 | BGModel = BGModel.to(device) |
| 46 | torch.cuda.empty_cache() |
| 47 | |
| 48 | |
| 49 | lq_paths = get_image_paths(input_path) |
| 50 | if len(lq_paths) ==0: |
| 51 | exit('No Image in the LR path.') |
| 52 | |
| 53 | |
| 54 | WEncoderPath='./checkpoints/net_w_encoder_860000.pth' |
| 55 | PriorModelPath='./checkpoints/net_prior_860000.pth' |
| 56 | SRModelPath='./checkpoints/net_sr_860000.pth' |
| 57 | YoloPath = './checkpoints/yolo11m_short_character.pt' |
| 58 | |
| 59 | TextModel = MARCONetPlus(WEncoderPath, PriorModelPath, SRModelPath, YoloPath, device=device) |
| 60 | |
| 61 | print('{:>25s} : {:s}'.format('Model Name', 'MARCONetPlusPlus')) |
| 62 | if use_cuda: |
| 63 | print('{:>25s} : {:<d}'.format('GPU ID', torch.cuda.current_device())) |
| 64 | else: |
| 65 | print('{:>25s} : {:s}'.format('GPU ID', 'No GPU is available. Use CPU instead.')) |
| 66 | torch.cuda.empty_cache() |
| 67 | |
| 68 | L_path = input_path |
| 69 | E_path = output_path # save path |
| 70 | print('{:>25s} : {:s}'.format('Input Path', L_path)) |
| 71 | print('{:>25s} : {:s}'.format('Output Path', E_path)) |
| 72 | if aligned: |
no test coverage detected