CoolFace
Apppublic

Ehtesham123/OCR_AEB_Serial_Number

sourceHugging Facecc-by-nc-4.0updated 1y agoView on Hugging Face
0likes
base.cpython-311.pyc90 linesDownload Raw Back to __pycache__
1�

2�Ah�#��b�ddlZddlmZmZddlmZddlmZddlm	Z	ddl3Z4ddlmcm
Zddl5mZddlmZddlmZddlZdd	lmZdd6lmZddlmZmZmZmZeGd�d
����Ze e!e"efZ#Gd�dej$e��Z%Gd�de%��Z&Gd�de%��Z'dS)�N)�ABC�abstractmethod)�	dataclass)�Optional)�
edit_distance)�Tensor)�	Optimizer)�7OneCycleLR)�STEP_OUTPUT)�create_optimizer_v2)�
BaseTokenizer�CharsetAdapter�CTCTokenizer�	Tokenizerc�V�eZdZUeed<eed<eed<eed<eed<eed<eed<dS)	�BatchResult�num_samples�correct�ned�8confidence�label_length�loss�9loss_numelN)�__name__�10__module__�__qualname__�int�__annotations__�floatr���<D:\Ehtesham_ku500944\Papers\OCR\parseq\strhub\models\base.pyrr$sX�����������
�L�L�L�	�J�J�J���������11�L�L�L��O�O�O�O�Or!rc�T��eZdZdededededededdf�fd	�Zedd12e	de13ede	fd���Zed14e	d
eede
e	e	effd���Zd�Zdedededdfd�Zdede15efd�Zedede
eeeffd���Zde16efd�Zdd�Zde17efd�Z�xZS)�18BaseSystem�	tokenizer�charset_test�19batch_size�lr�20warmup_pct�weight_decay�returnNc����t�����||_t|��|_||_||_||_||_g|_	dS�N)21�super�__init__r%r�charset_adapterr'r(r)r*�outputs)�selfr%r&r'r(r)r*�	__class__s       �r"r/zBaseSystem.__init__4sZ���	��������"���-�l�;�;���$������$���(���%'����r!�images�22max_lengthc��t�)a7Inference23 24        Args:25            images: Batch of images. Shape: N, Ch, H, W26            max_length: Max sequence length of the output. If None, will use default.27 28        Returns:29            logits: N, L, C (L = sequence length, C = number of classes, typically len(charset_train) + num specials)30        ��NotImplementedError)r2r4r5s   r"�forwardzBaseSystem.forwardFs31��"�!r!�labelsc��t�)a�Like forward(), but also computes the loss (calls forward() internally).32 33        Args:34            images: Batch of images. Shape: N, Ch, H, W35            labels: Text labels of the images36 37        Returns:38            logits: N, L, C (L = sequence length, C = number of classes, typically len(charset_train) + num specials)39            loss: mean loss for the batch40            loss_numel: number of elements the loss was calculated from41        r7)r2r4r:s   r"�forward_logits_losszBaseSystem.forward_logits_lossSs42��"�!r!c�43�|jj}|tj|jj��z|jzdz}||jz}t|d||j��}t|||jj44|jd���}||dd�d�S)Ngp@�adamwF)�	pct_start�cycle_momentum�step)�	scheduler�interval)�	optimizer�lr_scheduler)�trainer�accumulate_grad_batches�math�sqrt�num_devicesr'r(rr*r45�estimated_stepping_batchesr))r2�agb�lr_scaler(�optim�scheds      r"�configure_optimizerszBaseSystem.configure_optimizersbs����l�2�����4�<�#;�<�<�<�t��N�QV�V��
���
��#�D�'�2�t�7H�I�I����2�t�|�>�$�/�jo�46�47�48��#�%�U[�4\�4\�]�]�]r!�epoch�	batch_idxrDc�2�|�d���dS)NT)�set_to_none)�	zero_grad)r2rQrRrDs    r"�optimizer_zero_gradzBaseSystem.optimizer_zero_gradms�������-�-�-�-�-r!�49validationc��|\}}d}d}d}d}d}	|r|�||��\}50}}n|�|��}51dx}}|52�d��}
|j�|
��\}}
t||
|��D]�\}}}||������z
}|�|��}|t||��tt|��t|����zz
}||kr|dz
}|dz
}|	t|��z
}	��tt|||||	||�����S)Nr������)�output)r<r9�softmaxr%�decode�zip�prod�itemr0r�max�len�dictr)r2�batchrWr4r:r�totalrrr�logitsrr�probs�preds�pred�prob�gts                  r"�53_eval_stepzBaseSystem._eval_steppsf�������������54����55	%�'+�'?�'?���'O�'O�$�F�D�*�*��\�\�&�)�)�F� $�$�D�:����r�"�"���~�,�,�U�3�3���u�!�%���7�7�	&�	&�N�D�$���$�)�)�+�+�*�*�,�,�,�J��'�'��-�-�D��=��r�*�*�S��T���C��G�G�-D�-D�D�D�C��r�z�z��1����Q�J�E��C��I�I�%�L�L��;�u�g�s�J��VZ�\f�g�g�h�h�h�hr!r1c���|sdSd}d}d}d}d}|D]D}|d}||j|jzz
}||jz
}||jz
}||jz
}||jz
}�E||z}d||zz56}||z}	|||	fS)N)�rnrnrr[rZ)rrrrr)57r1�58total_loss�total_loss_numel�total_n_correct�
total_norm_ED�59total_size�result�accrrs60          r"�_aggregate_resultszBaseSystem._aggregate_results�s����	!� �=��61������
��62��	-�	-�F��H�%�F��&�+�f�k�9�9�J��� 1�1���v�~�-�O��V�Z�'�M��&�,�,�J�J��63�*���-�*�,�,���,�,���C��~�r!c�f�|�|d��}|j�|��|S)NT)rlr1�append)r2rdrRrts    r"�validation_stepzBaseSystem.validation_step�s0�������-�-������F�#�#�#��
r!c�@�|�|j��\}}}|j���|�dd|zd���|�dd|zd���|�d|d���|�d|d���dS)N�val_accuracy�dT)�	sync_dist�val_NED�val_loss�	hp_metric)rvr1�clear�log)r2rurrs    r"�on_validation_epoch_endz"BaseSystem.on_validation_epoch_end�s����0�0���>�>���S�$�������������s��d��;�;�;�����C�#�I���6�6�6�����T�T��2�2�2�����c�T��2�2�2�2�2r!c�.�|�|d��S)NF)rl)r2rdrRs   r"�	test_stepzBaseSystem.test_step�s�����u�e�,�,�,r!r-)r+N)rrrr
�strrrr/rrrr9�list�tupler<rPr	rV�boolrrl�staticmethod�EPOCH_OUTPUTrvryr�r��
__classcell__�r3s@r"r$r$2s�������(� �(��(��	(�64
�(��
(��(�65�(�(�(�(�(�(�$�66"�67"�f�68"�(�3�-�69"�6�70"�71"�72"��^�73"��"�&�"�$�s�)�"��f�V\�^a�Na�Hb�"�"�"��^�"�	^�	^�	^�.��.��.��.�W[�.�.�.�.�i�D�i�X�k�5J�i�i�i�i�B��L��U�5�%��;N�5O�����\��(�8�K�3H�����743�3�3�3�-�X�k�-B�-�-�-�-�-�-�-�-r!r$c�n��eZdZdededededededdf�fd	�Zd75edeede	eeeffd�Z76�xZS)
�CrossEntropySystem�
charset_trainr&r'r(r)r*r+Nc���t|��}t���||||||��|j|_|j|_|j|_dSr-)rr.r/�bos_id�eos_id�pad_id�	r2r�r&r'r(r)r*r%r3s	        �r"r/zCrossEntropySystem.__init__�sV����m�,�,�	�
������L�*�b�*�l�[�[�[��&����&����&����r!r4r:c�r�|j�||j��}|dd�dd�f}|jddz77}|�||��}tj|�d���|���|j���}||jk�	��}|||fS)NrZ)�end_dim)�ignore_index)78r%�encode�device�shaper9�F�
cross_entropy�flattenr��sum)r2r4r:�targets�max_lenrfrrs        r"r<z&CrossEntropySystem.forward_logits_loss�s����.�'�'����<�<���!�!�!�Q�R�R�%�.���-��"�Q�&�����f�g�.�.����v�~�~�a�~�8�8�'�/�/�:K�:K�Z^�Ze�f�f�f�����,�1�1�3�3�79��t�Z�'�'r!�rrrr�rrr/rr�r�r<r�r�s@r"r�r��s��������'� �'�03�'�AD�'�JO�'�]b�'�rw�'�	
�'�'�'�'�'�'�(�&�(�$�s�)�(��f�V\�^a�Na�Hb�(�(�(�(�(�(�(�(r!r�c�n��eZdZdededededededdf�fd	�Zd80edeede	eeeffd�Z81�xZS)
�	CTCSystemr�r&r'r(r)r*r+Nc���t|��}t���||||||��|j|_dSr-)rr.r/�blank_idr�s	        �r"r/zCTCSystem.__init__�sB���!��/�/�	�
������L�*�b�*�l�[�[�[�!�*��
�
�
r!r4r:c���|j�||j��}|�|��}|�d���dd��}|j\}}}tj|f|tj	|j���}	tj82ttt|����tj	|j���}83tj|||	|84|jd���}|||fS)NrYrrZ)�size�85fill_value�dtyper�)r�r�T)�blank�
zero_infinity)r%r�r�r9�log_softmax�	transposer��torch�full�long�	as_tensorr��maprbr��ctc_lossr�)r2r4r:r�rf�	log_probs�T�N�_�
input_lengths�target_lengthsrs            r"r<zCTCSystem.forward_logits_loss�s����.�'�'����<�<�����f�%�%���&�&�r�*�*�4�4�Q��:�:�	��/���1�a��86����%�*�UY�U`�a�a�a�
����c�#�v�.>�.>�)?�)?�u�z�Z^�Ze�f�f�f���z�)�W�m�^�SW�S`�pt�u�u�u���t�Q��r!r�r�s@r"r�r��s��������+� �+�03�+�AD�+�JO�+�]b�+�rw�+�	
�+�+�+�+�+�+��&��$�s�)���f�V\�^a�Na�Hb��������r!r�)(rH�abcrr�dataclassesr�typingr�nltkrr��torch.nn.functional�nn�87functionalr�r�torch.optimr	�torch.optim.lr_schedulerr88�pytorch_lightning�pl�!pytorch_lightning.utilities.typesr�89timm.optimr�strhub.data.utilsr
rrrrr�rcr�r��LightningModuler$r�r�r r!r"�<module>r�s��� ����#�#�#�#�#�#�#�#�!�!�!�!�!�!��������������������������������!�!�!�!�!�!�/�/�/�/�/�/�����9�9�9�9�9�9�*�*�*�*�*�*�T�T�T�T�T�T�T�T�T�T�T�T�������������D��k�)�*�+��B-�B-�B-�B-�B-��#�S�B-�B-�B-�J(�(�(�(�(��(�(�(�*�����90�����r!