# 이미지 불러와서 tensor로 변환 및 사이즈 조정
simple_transform = transforms.Compose([transforms.Resize((224,224))
                                       ,transforms.ToTensor()
                                       ,transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                      ])
img = Image.open('../image/cat_1.jpg')
img = simple_transform(img)

# unsqueeze_를 통해 shape에 (1,3,224,224) 로 변경 --> batch_size를 추가한것 --> 모델에 넣기위해
img.unsqueeze_(0)

# pretrained 모델을 불러와서 사용한다. --> 직접 train시킨 모델도 상관없지만 현재 목표는 모델 학습이 아니기 때문에 pretrained모델 사용
vgg = models.vgg16(pretrained=True).cuda()

class LayerActivations():
	# layer의 output을 저장할 공간을 만들자.
    features=None
    
	# hook을 정의하고 정한 layer_num에 hook을 등록할 수 있게 만든다.
    def __init__(self,model,layer_num):
        self.hook = model[layer_num].register_forward_hook(self.hook_fn)
    
	# 정해진 layer에서 forward함수가 실행되면 output을 저장할 수 있는 함수를 만든다.
    def hook_fn(self,module,input,output):
        self.features = output.cpu().data.numpy()
		
	# 등록된 hook을 없애줄 수 있는 함수를 만든다.
    def remove(self):
        self.hook.remove()

# vgg.features의 첫번째(0) 에 hook을 등록하고 output을 conv_out에 받도록 한다.
conv_out = LayerActivations(vgg.features,0)

# 가장 처음 가져온 이미지를 vgg모델에 넣어보자.
# 이미지가 모델을 통과할때 위에서 정의한 conv_out에 첫번째 conv layer를 통과하며 내놓은 output을 저장했을 것이다.
o = vgg(Variable(img.cuda()))

# 등록된 hook을 제거해준다.
conv_out.remove()

# conv_out객체의 features값을 가져오면 우리가 알고싶은 첫번째 conv layer를 통과했을때의 output을 알 수 있다.
act = conv_out.features

# 30개의 채널만을 plot으로 그려보자.
fig = plt.figure(figsize=(20,50))
fig.subplots_adjust(left=0,right=1,bottom=0,top=0.8,hspace=0,wspace=0.2)
for i in range(30):
    ax = fig.add_subplot(12,5,i+1,xticks=[],yticks=[])
    ax.imshow(act[0][i])

# 마지막 conv layer는 28번째 feature이다.
# 모든 feature를 확인해보고자 한다면 vgg.features 를 입력해서 확인해보자.
conv_out = LayerActivations(vgg.features,28)

o = vgg(Variable(img.cuda()))

conv_out.remove()

act = conv_out.features

fig = plt.figure(figsize=(20,50))
fig.subplots_adjust(left=0,right=1,bottom=0,top=0.8,hspace=0,wspace=0.2)
for i in range(30):
    ax = fig.add_subplot(12,5,i+1,xticks=[],yticks=[])
    ax.imshow(act[0][i])