문제 정의:

목표는 calibrated 된 카메라를 이용해 촬영된 인물의 head pose와 추정된 head pose의 uncertainty를 world coordinate frame에서 모델링 하는 것을 목표로 한다. 

+ 추가로 head pose(=face pose)를 이용해 인물의 시선을 예측하고 예측된 gaze point의 불확실성을 구하는 것까지 해본다.

 

노테이션 정의:

Pose $X^{frame}_{object}=(\psi, \phi, \theta, x, y, z) \in R^6$는  기준 coordinate frame 상에서 object의 pose 를 나타낸다. 

$$\psi : z axis 를 회전 축으로 하는 회전 각도, 단위:degree. $$

$$\phi : y axis 를 회전 축으로 하는 회전 각도, 단위: degree. $$

$$\theta : x axis 를 회전 축으로 하는 회전 각도, 단위: degree. $$

$$ x : object 의 x 좌표 $$

$$ y: object 의 y 좌표 $$

$$ z: object 의 z 좌표 $$

 

단, face coordinate frame 에서 face의 pose는 $X^{face}_{face}=(0,0,0,0,0,0)$ 으로 정의한다.

이 문제에서는 $X^{face}_{face}, X^{face}_{camera}, X^{camera}_{world} $ 3개의 pose 가 등장한다.

 

Pose 와 transformation matrix 관계

$$^{object}_{frame}H = f(X^{frame}_{object}) = \begin{bmatrix}^{object}_{frame}R&-^{object}_{frame}RT_{object}\\0&1\\ \end{bmatrix}$$

$$^{object}_{frame}R=Rz(\psi)Ry(\phi)Rx(\theta),  T_{object}=\begin{bmatrix} x\\y\\z \end{bmatrix}$$

$$p^{object} = ^{object}_{frame}H \times p^{frame}, p^{frame}=\begin{bmatrix} x\\y\\z\\1 \end{bmatrix}$$

$$^{world}_{face}H = ^{world}_{camera}H \times ^{camera}_{face}H$$

$$X^{frame}_{object}=f^{-1}(^{object}_{frame}H)$$

 

좌표계 정의

 

Pose 추정

$X^{face}_{camera} = SolvePnP(q^{face}_{i}, p^{image}_{landmark-i},K, dist)$

$X^{camera}_{world} = SolvePnP(q^{world}_{i}, p^{image}_{target-i}, K, dist)$

 

K: projection matrix, dist = distortion parameter

 

$p^{image}_{landmark-i}$ : image 좌표계 상의 i번째 face landmark 점 좌표 (x,y). 

$q^{face}_{i}$ : face coordinate frame 상의 $ p^{image}_{landmark-i}$ 대응점 (x,y,z).

$p^{image}_{target-i}$ : image 좌표계 상의 i번째 calibration target landmark 점 좌표 (x,y). 

$q^{world}_{i}$ : world coordinate 상의  $p^{image}_{target-i}$ 대응점 (x,y,z).

 

추정 오차 공분산 계산:

$p^{image}_{i} = f(X^{face}_{camera}) \times q^{face}_{i}$ 이 수식을 테일러 시리즈를 이용해 선형 근사 하면

$ \overline{p}^{image}_{i}+\Delta p = f( \overline{X}^{face}_{camera}+\Delta X) \times q^{face}_{i} $

 

$\Delta p \approx \frac{\partial f}{\partial X} \Delta X^{face}_{camera} = M_{i} \Delta X^{face}_{camera}$

이 식을 least square 로 풀면 Pose $X^{face}_{camera}$ 의 추정 오차 공분산은 

$C_{x} = E[\Delta X \Delta X^{T}] = (M^{T}M)^{-1}M^{T}E(\Delta p \Delta p^{T})((M^{T}M)^{-1}M^{T})^{-1}$

$E(\Delta p \Delta p^{T}) : ^{camera}_{face}H 를 이용해 projection 한 point의 공분산 행렬$ 로 나타낼 수 있다.

$C_{x}$는 6x6 사이즈의 행렬이다.

 

World-camera 의 추정 오차 공분산 행렬도 $C_{w} = 6 \times 6$  이고  위와 동일 한 방식으로 calibration target을 이용해 구할수 있다.

 

최종 $^{world}_{face}H$의 추정 오차 공분산을 $C_{y}$ 라 하면

$X^{face}_{world}=f^{-1}(^{world}_{face}H)$

$^{world}_{face}H= ^{world}_{camera}H \times ^{camera}_{face}H = f(X^{camera}_{world}) \times f(X^{face}_{camera})$ 관계에 의해 공분산 전파(propagation) 식에 의해 아래와 같이 구할 수 있다. 

 

$C_{y} = J_{X} C_{x} J^{T}_{X} + J_{W} C_{W} J^{T}_{w}, C_{y} = 6 \times 6 행렬$

 

$J_{x} = \frac{\partial f^{-1}(^{world}_{face}H)}{\partial X^{face}_{camera}} $

 

$J_{w} = \frac{\partial f^{-1}(^{world}_{face}H)}{\partial X^{camera}_{world}} $

 

이를 이용해 face pose 는 world에서 $X^{face}_{world}=f^{-1}(^{world}_{face}H)$ 를 평균으로 하고 $C_{y}$ 분산으로 하는 gaussian pdf 를 따른다고 볼 수 있다. 

 

다음으로 head pose로 부터 gaze point를 예측 해보자. 

여기서는 시선의 방향이 얼굴의 전면부 즉 코끝이 가리키는 방향과 같다고 가정한다. 

 

gaze point $g$를 아래와 같이 정의하자. 

$$g=\begin{bmatrix} x_{g}\\y_{g}\\z_{g} \end{bmatrix} = p_{face} + tV_{gaze} = g(^{world}_{face}H)      --------(func gaze)$$

$V_{gaze} =\begin{bmatrix} g_{x}\\g_{y}\\g_{z} \end{bmatrix}: gaze direction vector, ^{world}_{face}R의 마지막 컬럼, 즉 ^{world}_{face}H[:,2] = f(X^{face}_{world})[:,2]$

$p_{face} = \begin{bmatrix} x_{face}\\y_{face}\\z_{face} \end{bmatrix}$, 얼굴 위의 한점 여기선 face coordinate frame origin의  world coordinate frame 상의 좌표로 설정, 즉 $^{world}_{face}H[:,3]=f(X^{face}_{world})[:,3]$

 

이때 $V_{gaze}$는 시선의 방향을 나타네는 벡터 즉 시선의 방향 벡터라고 볼 수 있다. 가정에 의해 시선의 방향은 얼굴 평면과 수직(perpendicular) 이므로 world 좌표계에서 head pose를 나타내는 rotation matrix의 3번째 컬럼 즉 face coordinate frame의 z axis에 해당한다. 3차원 공간상에서 방향벡터 $V$와 평행하고 사람의 얼굴 위의 한점(눈 사이의 한점을 잡는게 가장 좋으나 여기서는 코끝으로 가정했다.)을 지나는 직선을 구하는게 목적이므로 $p_{face}$는 코끝의 world coordinate frame 상의 좌표로 가정하자. 

 

이렇게 하면 world coordinate frame 상의 z-x평면위에서 이미지 상의 특정 인물이 바라보고 있는 좌표(gaze point)와 불확실성은 아래와 같이 구할 수 있다.

$t= -\frac{p_{face}}{g_{y}} $ 로 설정 하면 y=0이 되므로 world 좌표계 상에서 x,z평면과 만나는 gaze point를 구할 수 있고

gaze point 와 pose $X^{face}_{world}$의 관계에 따라 gaze point 의 공분산은 

$C_{g} = J_{g} C_{y} J^{T}_{g} , 3 \times 3 행렬$

$J_{g} = \frac {\partial g(^{world}_{face}H)}{\partial g}, 3 \times 6 행렬$

로 구할 수 있다. 

 

 

다음은 face pose(=head pose) estimation에 사용된 landmark를 표시한 그림이다. 

 

'Deeplearning > toyproject' 카테고리의 다른 글

[Deskew for ocr] Rotation correction v2  (0) 2022.03.06
[Deskew for ocr] Rotation correction  (0) 2022.02.15

이전 글 [Deskew for ocr] Rotation correction 에서 regression으로 회전된 문서를 바로잡는 방식을 시도해봤다.

 

이전 글에서 쓴 모델의 문제는 크게 2가지

1. 회전되지 않는 문서를 회전된 문서로 오인식해 오히려 이상하게 회전시키는 문제가 발생한다.

2. 정확도가 굉장히 떨어진다.

 

이를 해결하기 위해 새로운 방식을 조사, 적용 했고 결과가 좋아서 공유한다.

 

이번엔 회전된 문서를 바로잡는 문제를 분류 문제로 정의 하고 해결해 본다. 이 방식은 ocropus3 를 참고 했다.

 

코드는 아래 링크의 devStream_fft branch 에 있다.

link: https://github.com/pajamacoders/ocrDeskew/tree/devStream_fft

문제 정의

문서가 얼마나 회전되어있는 지를 분류 문제로 정의 하고 풀기 위해서 각 회전의 정도에 class를 부여 해야 한다.

 

나는 0.5도 단위를 하나의 클래스로 정의했다.

예를 들자면 아래와 같은 방식이다. 문서가 회전된 각도를 degree 라고 표현 했을 때 회전의 정도(degree) 가 -1 도에서 -0.5 도 이내이면 class 0에 배정 하는 방식이다.

range -1< degree<-0.5 -0.5<= degree <0 0<= degree <0.5
class 0 1 2

이러한 방식으로 -89~89도 사이에서 회전된 문서의 rotation correction 문제는

356 클래스를 가지는 분류 문제로 정의 할 수 있다.

 

개발 환경

개발 환경은 ngc repo에서 아래 이미지를 다운 받았다.

docker image : nvcr.io/nvidia/pytorch:22.01-py3

train metric tracking: mlflow

 

전처리

전처리 과정은 이전 글 [Deskew for ocr] Rotation correction 의 전처리와 거의 유사하지만 rotation의 정도에 class를 대응 하는 부분만 차이가 있다.

 

이전 글에서 바뀐 GT 생성 부분인 RandomRotation class 는 아래와 같다.

 

class RandomRotation(object):
    def __init__(self, ratio, degree, buckets=None):
        self.variant = eval(degree) if isinstance(degree, str) else degree
        self.ratio = eval(ratio) if isinstance(ratio, str) else ratio
        self.buckets = eval(buckets) if isinstance(buckets, str) else buckets

    def __call__(self, inp):
        if  np.random.rand()<self.ratio:
            deg = np.random.uniform(-self.variant, self.variant-0.1)
            img = inp['img']
            h,w= img.shape
            matrix = cv2.getRotationMatrix2D((w/2, h/2), deg, 1)
            dst = cv2.warpAffine(img, matrix, (w, h),borderValue=0)
            inp['img'] = dst
        else:
            deg = 0

        if self.buckets:
            rad = np.deg2rad(deg)
            range_rad = np.deg2rad(self.variant)
            bucket = int(self.buckets * (rad+range_rad) / (2*range_rad))
            inp['rot_id'] = bucket # 이 값이 문서가 회전된 정도를 class에 할당 한 값 즉 target class 값이다.

        inp['degree'] = deg
        return inp

 

모델

이번 모델의 특이점은 중간에 fft를 사용 하는 layer 가 들어간다는 것이다.

 

모델은 아래와 같다.

class DeskewNetV4(nn.Module):
    def __init__(self, buckets, last_fc_in_ch, pretrained=None):
        super(DeskewNetV4, self).__init__()
        buckets = eval(buckets) if isinstance(buckets, str) else buckets
        assert isinstance(buckets, int), 'buckets must be type int'
        k=5
        self.block1 = nn.Sequential(
                ConvBnRelu(1,8,k,padding=k//2),
                nn.MaxPool2d(2,2), #256x256
                ConvBnRelu(8,16,k,padding=k//2),
                nn.MaxPool2d(2,2), #128x128
            )
        
        self.block2 = ConvBnRelu(16,8,k,padding=k//2)
        self.fc = nn.Sequential(nn.Linear(131072,last_fc_in_ch, bias=False),
            nn.BatchNorm1d(last_fc_in_ch),
            nn.ReLU(True),
            nn.Linear(last_fc_in_ch, buckets, bias=False))

        self.__init_weight()
        if pretrained:
            self.load_weight(pretrained)

    def forward(self, x):
        out = self.block1(x)
        out = torch.fft.fft2(out)
        out = out.real**2+out.imag**2 
        out = torch.log(1.0+out)
        out = self.block2(out)
        bs,c,h,w = out.shape
        out = out.reshape(bs,-1)
        out = self.fc(out)

        return out

 

결과

모델의 last_fc_in_ch의 값으로 128을 사용 할 경우 아래와 같은 결과를 얻었다.

학습은 총 800에폭을 돌렸는데 굳이 이럴 필요까진 없었다.

optimizer로 adam을 사용 했고 lr_schedule은 cosineannealing을 lr range  1e-3 ~1e-6으로 사용 했다.

 

ce_loss:0.1241

precision:0.9466

recall:0.9426

f1_score:0.9427

 

아래는 train, validation 시의 f1 score의 값을 나타낸 그래프 이다.

200epoch 쯤 되면 성능 향상은 거의 없는 것을 볼 수 있다.

아래 그림은 임의의 숫자를 적은 문서로 테스트 한 결과 이다.

왼쪽이 입력으로 들어간 회전된 문서이고 오른쪽인 inference로 회전을 바로잡은 결과이다.

regression 모델 보다는 전체적인 결과가 훨씬 좋다.

 

- 끝 -

배경

이번 포스팅에서는 ocr 성능을 높이기 위해 간단한 전처리를 통해 회전된 문서를 올바르게 돌려 주는 문제를 풀어보고자 한다.

 

OCR 은 optical character recognition의 약자로 hand writing, 인쇄된 문서 등을 카메라로 찍거나 스캔했을때 그 문서 내의 글자를 인식해 문서를 전산화 할때 자주 사용되는 기술이다.

 

이때 ocr의 성능을 떨어뜨리는 문제 중 하나는 입력으로 들어오는 문서가 회전되어있는 경우 이다.

문서가 회전된 상태로 스캔되면 각 문자 자체는 인식이 할수 있지만 문자를 단어로 머지(merge)하는 과정이나 숫자의 연속인 여권번호, 운전면허 번호, 주민등록 번호 등 긴 문자의 경우 하나의 시퀀스로 머지 되지 않아 잘못된 패턴으로 인식 되는 경우가 종종 발생해  최종 인식률이 떨어질 수 있다. 

 

이런 문제를 해결하기 위해 인식 네트워크 자체가 문서 얼라인먼트 기능을 가지도록 설계할 수도 있고, hand craft feature extraction을 기가 막히게 설계 하고 이를 이용한 homography 계산 알고리즘을 설계하는 방식도 있겠으나!!!

여기서는 간단한 네트워크 설계하고 학습해 rotation correction하는 방식을 시도해 본다.

 

* 이 포스팅에 사용된 문서는 공공문서 양식 중 하나이고 이런 DB는 'aihub -> 음성/자연어->공공 행정문서 OCR'에서 아주 손쉽게 구할수 있다.

 

이 포스트에서 개발한 코드는 이 링크에 있다.

code :https://github.com/pajamacoders/ocrDeskew

목표

Fig 1. 왼쪽은 회전된 문서, 오른쪽은 회전 없이 정상적으로 스캔된 문서

이 포스팅의 목적을 명확히 나타내는 그림이 Fig 1. 이다. Fig 1 의 왼쪽은 회전된 문서를 보여준다. 우리의 목표는 이렇게 회전된 문서를 오른 쪽과 같이 회전이 없는 상태로 만드는 것이다.

 

문제 정의

나는 이 포스팅의 목표인 rotation correction 문제를 image orientation prediction 문제로 정의 했다.

 

사고의 흐름은 입력 이미지로 부터 이미지의 회전 정도(orientation)를 구하면 그 회전의 크기만큼 반대 방향으로 회전을 시켜 줌으로서 이미지를 정상적으로 만들수 있기 때문이다.

 

목표 추정 범위는 -30~30degree 이내의 회전으로 정했다.

회전의 크기를 일정 step으로 양자화 해서 각 구간에 class를 부과해 classification으로 해결 할 수도 있을것 같지만

회전 크기를 degree로 바로 추정하는 regression 문제로 정의 하고 풀고자 한다.

 

개발 환경

개발 환경은 ngc repo에서 아래 이미지를 다운 받았다.

docker image : nvcr.io/nvidia/pytorch:22.01-py3

train metric tracking: mlflow

 

전처리

전처리는 과정에서는 크게 아래 4가지를 수행했다.

  1. Color conversion(BGR 2 GRAY)
  2. Resize
  3. Rotation (이 과정은 target GT를 생성하는 과정이다.)
  4. Normalization

이때 resize시에 특정한 문제가 발생했고 이를 해결하기 위해 약간의 꼼수(=테크닉)가 사용되어 그 부분에 대해 잠시 설명 하고자 한다.

 

입력 이미지가 대부분 A4등의 문서를 스캔한 것이기 때문에 풀고자하는 문제에 비해 불필요 하게 크다고 생각했고 입력 이미지의 크기를 512x512로 정의 했다.

하지만 입력 중 어떤 이미지들은 3508x2480(width x height)의 크기를 지녔고 이걸 (512x512)로 리사이즈 하니 다음과 같은 이미지가 생성되었다.

Fig 2. 회전된 이미지

예제의 이미지는 그나마 좀 나은 편인데 심한 것들은 의미있는 글자나 선등이 sampling이 거의 안되어 노이즈 처럼 점만 뿌려져 있는 것처럼 보이기도 한다.

문제는 이러한 텍스트 문서의 경우 배경에 비해 정보를 담고 있는 글자나 테이블을 이루는 선등의 면적이 매우 작아 리사즈 할때 텍스트나 선 부분에서 샘플링되는 양이 매우 작기 때문이다.

 

위에 서 말했듯이 나는 원인이 정보를 포함하는 텍스트나 선이 차지하는 면적이 작기 때문으로 보았기 때문에 리사이즈 하기 전에 이 정보를 포함하는 부분의 면적을 키우기로 했다.

 

가장 먼저 떠오른 것은 dilate 연산이다.

dilate을 적용하기 위해 이미지의 색을 반전 시켜 준다. 색 반전의 이유는 입력 이미지는 흰색의 배경을 가지는데 기본적으로 dilate연산은 intensity가 큰 부분은 확장하는 효과를 가지기 때문에 내가 면적을 키우고자 하는 텍스트와 선 등 정보를 가진 pixel이 큰intensity 가지고 background가 낮은 intensity를 가지게 하기 위해서 이다.

아래 코드와 같이 이미지를 읽고 gray scale로 색변환 한 후 background와 foreground의 색상을 반전 시켰다.

    img = cv2.imread(self.img_pathes[index])
    # transform
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    if self.inverse_color:
        img = 255-img # inverse image to apply dilate
    res_dict = {'img':img, 'imgpath':self.img_pathes[index]}

다음으로 적용한 테크닉은 3508 x 2480 을 곧바로 512x512로 리사이즈 하는게 아니라

dilate ->1/2 resize -> 1/2 resize-> resize to 512x512 로 리사이즈 스탭을 진행해 orientation 추출에 필요한 정보를 최대한 많이 보존 하게 하는 것이다.

이 부분의 코드는 다음과 같다.

class Resize(object):
    def __init__(self, scale=4):
        assert (scale!=0) and (scale&(scale-1))==0, 'scale must be power of 2'
        self.iter = np.log2(scale).astype(int)
    
    def __call__(self, inp):
        img = inp['img']
        h,w = img.shape
        inp['org_height']=h
        inp['org_width']=w
        k = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
        img = cv2.dilate(img, k)
        for i in range(self.iter):
            h, w = h//2, w//2
            img=cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC)
            h, w = img.shape
        inp['img']=cv2.resize(img, (512, 512), interpolation=cv2.INTER_CUBIC)
        return inp

이렇게 만들어진 이미지는 아래 Fig 3과 같다.  Fig 2와 텍스트와 선분의 정보 손실 정도를 비교해 보자.

리사이즈 방식에 따라 최종 목적인 regression 성능이 얼마나 달라지는지 ablation study를 진행해 보지는 않았지만 정보의 손실의 최소화 하고자 하는 측면에서는 이게 맞지 않나 싶다.

Fig 3. resize로 인한 글자 정보손실을 최소화한 이미지

또 다른 방식으로는 텍스트의 뭉개짐을 무시하고 높은 intensity를 가지는 부분을 최대화 하고자 한다면 아래와 같이 입력을 변환 하는 방식도 가능 할 것이다.

Fig 4.

Fig 4에서는 글자를 알아 보기는 힘들지만 이미지의 orientation을 결정하는데에 foreground 정보의 양이 중요할 경우 유용 할 것이다. Fig 4는 위 Resize 클래스 에서 cv2.dilate() 함수를 for 문 안에 cv2.resize함수 호출 바로 앞으로 옮긴 경우 만들어 지는 이미지 이다.

 

GT 생성

위 이미지 전처리 과정에서 Rotation 과정은 아래와 같은 이유에서 추가되었다.

 

이미지가 회전된 정도를 GT로 만들어 놓은 경우는 거의 없기 때문에 정상적인 이미지를 opencv를 이용해 랜덤으로 -30~30degree 사이 각도로 회전시켜 회전된 정도를 GT로 사용 한다.

 

아래코드는 resize된 이미지를 임의의 각도로 회전하고 회전한 각도를 target GT로 저장하는 코드이다.

 

class RandomRotation(object):
    def __init__(self,degree):
        self.variant = eval(degree) if isinstance(degree, str) else degree

    def __call__(self, inp):
        deg = np.random.uniform(-self.variant, self.variant)
        img = inp['img']
        h,w= img.shape
        matrix = cv2.getRotationMatrix2D((w/2, h/2), deg, 1)
        dst = cv2.warpAffine(img, matrix, (w, h),borderValue=0)
        inp['img'] = dst
        inp['degree'] = deg
        return inp

 

모델 설계

입력의 회전된 각도를 추정하는 아주 가벼운 네트워크를 아래와 같이 구성했다.

import torch
import torch.nn as nn


class ConvBnRelu(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, dilation=1,
                 groups=1):
        super(ConvBnRelu, self).__init__()
        self.conv_bn_relu = nn.Sequential(
            nn.Conv2d(in_channel, out_channel, kernel_size, stride, padding, dilation, groups,
                      False),
            nn.BatchNorm2d(out_channel),
            nn.ReLU(True))

    def forward(self, x):
        return self.conv_bn_relu(x)

class DeskewNet(nn.Module):
    def __init__(self, pretrained=None):
        super(DeskewNet, self).__init__()
        k=3
        self.backbone = nn.Sequential(
            ConvBnRelu(1,8,k,padding=k//2),
            nn.MaxPool2d(2,2), #256x256
            ConvBnRelu(8,16,k,padding=k//2),
            nn.MaxPool2d(2,2), #128x128
            ConvBnRelu(16,32,k,padding=k//2),
            nn.MaxPool2d(2,2), #64x64
            ConvBnRelu(32,64,k,padding=k//2),
            nn.MaxPool2d(2,2), #32x32
            ConvBnRelu(64,64,k,padding=k//2),
            nn.MaxPool2d(2,2), #16x16
        )
        self.avgpool = nn.AvgPool2d((16,16))
        self.fc = nn.Sequential(nn.Linear(64,64),nn.Linear(64,1))
        self.__init_weight()
        if pretrained:
            self.load_weight(pretrained)

    def forward(self, x):
        out = self.backbone(x)
        out = self.avgpool(out)
        out = self.fc(out.squeeze())
        return out

    def __init_weight(self):
        for layer in self.modules():
            if isinstance(layer, nn.Conv2d):
                nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')
            else: 
                pass

 

학습

학습 환경은 아래와 같이 구성했다.

optimizer: Adam

lr_scheduler: CosineAnnealingLR -> 초기 lr =0.001, eta_min 1e-6, T_Max=300

loss: MSELoss

batch: 128

아래는 학습 코드 메인 문이다. train, valid 함수 구현은 뻔하니 생략한다.

if __name__ == "__main__":
    args = parse_args()
    with open(args.config, 'r') as f:
        cfg = json.load(f)
        cfg['config_file']=args.config
        if args.run_name:
            cfg['mllogger_cfg']['run_name']=args.run_name
    
    tr = build_transformer(cfg['transform_cfg'])
    train_loader, valid_loader = build_dataloader(**cfg['dataset_cfg'], augment_fn=tr)

    logger.info('create model')
    model = build_model(**cfg['model_cfg'])#torch.hub.load('pytorch/vision:v0.10.0', cfg['model_cfg']['type'])
    model.cuda()
    logger.info('create loss function')
    fn_loss = torch.nn.MSELoss()

    logger.info('create optimizer')
    opt=torch.optim.Adam(model.parameters(), **cfg['optimizer_cfg']['args'])
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt,**cfg['lr_scheduler_cfg']['args'])

    max_epoch = cfg['train_cfg']['max_epoch']
    valid_ecpoh = cfg['train_cfg']['validation_every_n_epoch']
    logger.info(f'max_epoch :{max_epoch}')
    logger.info('set mlflow tracking')
    mltracker = MLLogger(cfg, logger)
    for step in range(max_epoch):
        train(model, train_loader, fn_loss, opt, mltracker, step)
        if (step+1)%valid_ecpoh==0:
            valid(model, valid_loader, fn_loss,  mltracker, step)
        lr_scheduler.step()

 

결과

아래는 train, validation loss를 mlflow tracking 으로 추적한 결과 이다. loss가 초반에 급격히 떨어지고 수렴은 하지만

validation이 중간 중간 불안정 하게 튀는 모습을 볼 수 있다.  이건 해결을 좀 해야겠다.

Fig5. loss graph

아래 Fig 6.의 왼쪽은 회전왼 입력 이미지이고 오른쪽 그림은 orientation을 추정해 correction 한 이미지 이다. 정확하진 않지만 생태는 많이 나아 졌다. 좀더 정확하게 correction 되도록 data검증 및 네트우크 수정을 해봐야겠다.

 

끝.

+ Recent posts