test_ocr_1 / test.py
B2W1234's picture
Upload 53 files
0a96b76 verified
"""Modify to Remove Argument Parser"""
"""
Copyright (c) 2019-present NAVER Corp.
MIT License
"""
# -*- coding: utf-8 -*-
import sys
import os
import time
import argparse
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
from PIL import Image
import cv2
from skimage import io
import numpy as np
import craft_utils
import imgproc
import file_utils
import json
import zipfile
from craft import CRAFT
from collections import OrderedDict
def copyStateDict(state_dict):
if list(state_dict.keys())[0].startswith("module"):
start_idx = 1
else:
start_idx = 0
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = ".".join(k.split(".")[start_idx:])
new_state_dict[name] = v
return new_state_dict
def test_net(net, image, text_threshold, link_threshold, low_text, cuda, poly, args, refine_net=None):
t0 = time.time()
# resize
# cứ biết là lấy ra kích thước ảnh mới , tỉ lệ với chiều rộng chiều cao ảnh, kích thước bản đồ nhiệt
img_resized, target_ratio, size_heatmap = imgproc.resize_aspect_ratio(image, args.canvas_size,
interpolation=cv2.INTER_LINEAR,
mag_ratio=args.mag_ratio)
# đây chính là lấy tỉ lệ với chiều rộng và chiều cao ảnh còn là gì thì xuống xem
ratio_h = ratio_w = 1 / target_ratio
# preprocessing Bước tiền xử lý
x = imgproc.normalizeMeanVariance(img_resized)
'''
chú ý rằng tensor trong pytorch bằng với array trong numpy
trong 2 lệnh tiếp theo
đầu tiên là thay đổi thứ tự và dòng tiếp tăng số chiều của nó lên
mục đích chưa rõ
https://github.com/pytorch/pytorch/issues/44541
'''
x = torch.from_numpy(x).permute(2, 0, 1) # [h, w, c] to [c, h, w]
x = Variable(x.unsqueeze(0)) # [c, h, w] to [b, c, h, w]
if cuda:
x = x.cpu()
# forward pass
# từ net(x) lấy ra y và feature nữa, mục đích chưa rõ
with torch.no_grad():
y, feature = net(x)
# make score and link map
score_text = y[0, :, :, 0].cpu().data.numpy()
score_link = y[0, :, :, 1].cpu().data.numpy()
# refine link
if refine_net is not None:
with torch.no_grad():
y_refiner = refine_net(y, feature)
score_link = y_refiner[0, :, :, 0].cpu().data.numpy()
t0 = time.time() - t0
t1 = time.time()
# những tham số truyền vào này này đều là tham số default từ file với pipeline
# mục đích là lấy ra boxes, polys, det_scores mới
# quá trình xử lý chưa rõ
# Post-processing gọi hàm từ file craft_utils
# boxes để lấy tọa độ, polys có vẻ là dương tự nhưng mà dưới dạng weight hoặc cách implement khác
#
boxes, polys, det_scores = craft_utils.getDetBoxes(score_text, score_link, text_threshold, link_threshold, low_text,
poly)
# coordinate adjustment gọi hàm từ file craft_utils
boxes = craft_utils.adjustResultCoordinates(boxes, ratio_w, ratio_h)
polys = craft_utils.adjustResultCoordinates(polys, ratio_w, ratio_h)
for k in range(len(polys)):
if polys[k] is None: polys[k] = boxes[k]
t1 = time.time() - t1
# render results (optional)
render_img = score_text.copy()
render_img = np.hstack((render_img, score_link))
ret_score_text = imgproc.cvt2HeatmapImg(render_img)
if args.show_time: print("\ninfer/postproc time : {:.3f}/{:.3f}".format(t0, t1))
return boxes, polys, ret_score_text, det_scores