Hugging Face's logo

Spaces:
Suchinthana
/
ColorGradingAE
Sleeping

App Files Community
ColorGradingAE
File size: 1,040 Bytes
fcd8868
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import cv2
import torch
import numpy as np

from trainer import (
    extract_features,
    Encoder,
    Decoder,
    DEVICE,
    INPUT_SIZE
)



encoder=Encoder().to(DEVICE)
decoder=Decoder().to(DEVICE)


encoder.load_state_dict(
    torch.load("encoder.pt")
)

decoder.load_state_dict(
    torch.load("decoder.pt")
)


encoder.eval()
decoder.eval()



def generate_grade(
        image_path,
        warmth=0):


    f=extract_features(
        image_path
    )


    x=torch.tensor(
        f
    ).float().to(DEVICE)


    with torch.no_grad():

        z=encoder(
            x.unsqueeze(0)
        )

        result=decoder(z)


    grade=result.cpu().numpy()[0]


    return grade



def apply_grade(
        image_path,
        grade):


    img=cv2.imread(
        image_path
    )

    hsv=cv2.cvtColor(
        img,
        cv2.COLOR_BGR2HSV
    )


    hsv[:,:,1]=np.clip(
        hsv[:,:,1]
        *
        (1+grade.mean()),
        0,
        255
    )


    return cv2.cvtColor(
        hsv,
        cv2.COLOR_HSV2BGR
    )