frozencherry/Forgery-Localization-App
1
1a
2 j}�h� � @ sN d dl Zd dlmZ G dd� dej�ZG dd� dej�ZG dd� dej�ZdS ) � N)�layersc s. e Zd Zd� fdd� Zdd� Zdd� Z� ZS ) �ChannelAttention� c s t t| �jf i |�� || _d S �N)�superr �__init__�ratio)�selfr �kwargs�� __class__� �?/home/ubuntu/Documents/forgery-localization-app/custom_layer.pyr s zChannelAttention.__init__c C s8 |d }t j|| j dddd�| _t j|ddd�| _d S )N������relu� he_normalT)�3activation�kernel_initializer�use_bias)r r )r �Denser �shared_dense_one�shared_dense_two)r � input_shape�channelsr
r
r �build4 s zChannelAttention.buildc C s� t �� |�}t �dd|jd f�|�}| �|�}| �|�}t �� |�}t �dd|jd f�|�}| �|�}| �|�}tj� || �}t �5� ||g�S )N� ) r �GlobalAveragePooling2D�Reshape�shaper r �GlobalMaxPooling2D�tf�nn�sigmoid�Multiply)r �inputs�avg_pool�max_pool� cbam_featurer
r
r �call s 67 89zChannelAttention.call)r ��__name__�10__module__� __qualname__r r r( �
__classcell__r
r
r r r s r c s. e Zd Zd� fdd� Zdd� Zdd� Z� ZS ) �SpatialAttention� c s t t| �jf i |�� || _d S r )r r. r � kernel_size)r r0 r11 r r
r r s zSpatialAttention.__init__c C s" t jd| j| jfdddd�| _d S )Nr �samer F)�paddingr r )r �Conv2Dr0 �conv2d�r r r
r
r r s zSpatialAttention.buildc C s` t �dd� �|�}t �dd� �|�}t jdd�||g�}| �|�}t �d�|�}t �� ||g�S )Nc S s t j| ddd�S �N� T)�axis�keepdims)r � reduce_mean��xr
r
r �<lambda>$ � z'SpatialAttention.call.<locals>.<lambda>c S s t j| ddd�S r6 )r �12reduce_maxr; r
r
r r= % r> r7 )r8 r" )r �Lambda� Concatenater4 �13Activationr# )r r$ r% r& �concatr' r
r
r r( # s 14zSpatialAttention.call)r/ r) r
r
r r r. s r. c s, e Zd Z� fdd�Zdd� Zdd� Z� ZS )�CBAMc s6 t t| �jf i |�� d}d}t|�| _t|�| _d S )Nr r/ )r rD r r �channel_attentionr. �spatial_attention)r r15 r r0 r r
r r , s16 17z
CBAM.__init__c C s | � |�}| �|�}|S r )rE rF )r r$ r' r
r
r r( 3 s 1819z CBAM.callc C s d S r r
r5 r
r
r r 8 s z20CBAM.build)r* r+ r, r r( r r- r
r
r r rD + s rD )�21tensorflowr Ztensorflow.kerasr �Layerr r. rD r
r
r
r �<module> s