| 35 | return model |
| 36 | |
| 37 | def process_image(img_path, model, device, args): |
| 38 | imgname = osp.splitext(osp.basename(img_path))[0] |
| 39 | print('processing image: ', imgname) |
| 40 | img = cv2.imread(img_path, cv2.IMREAD_COLOR).astype(np.float32) / 255. |
| 41 | img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float() |
| 42 | img = img.unsqueeze(0).to(device) |
| 43 | mean = np.array([0.485, 0.456, 0.406]) |
| 44 | std = np.array([0.229, 0.224, 0.225]) |
| 45 | |
| 46 | normalize(img, mean, std, inplace=True) |
| 47 | with torch.no_grad(): |
| 48 | output = model(img) |
| 49 | output = normalize(output, -1 * mean / std, 1 / std) |
| 50 | output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy() |
| 51 | if output.ndim == 3: |
| 52 | output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0)) |
| 53 | output = (output * 255.0).round().astype(np.uint8) |
| 54 | cv2.imwrite(osp.join(args.output, f'{imgname}_{args.model}.png'), output) |
| 55 | |
| 56 | def main(): |
| 57 | parser = argparse.ArgumentParser() |