B

    dcY.                 @   s   d dl Z d dlZd dlZd dlmZ d dlZd dlmZ d dl	m	Z	 d dl
Zd dlm
Z
 d dlZd dlZd dlZG dd deZdd	 Zdd
dZdS )
    N)Dataset)
ToTensorV2)glob)OrderedDictc               @   s&   e Zd ZdddZdd Zdd ZdS )	
DoppioDatasetNc          
   C   s<  i }t  | _d||t| d }d||| _d||t| d | _tj	|d}t
j|}	t
|	}
W d Q R X |
d }|
d | _|
d | _g | _tj	|d}t
j|}	t
|	| _W d Q R X t| d | _x| jD ]}
|
d	krtjd
dgd|d
< d | jd< q|
dkr2d | jd< q|
dkrtjddgd|d< d | jd< d | jd< d | jd< dd tt| jd D | _q|
dkrd | jd< q|
dkrd | jd< qW t t| jd tt| jd | _d}d| jkr|d}|d krHt t|tt|| jd< i | _x\t|D ]\}}|| j| j| < q(W n4i | _x,t| jd D ]\}}|| j| j| < q^W |d d }x|D ]}tj	||}t
j|}	t
|	}W d Q R X d}x0|d D ]$}|d }|| jd krd}P qW |dks|dkr| j| qW tj |f|| _!d | jd < d S )!Nz3http://222.231.59.68/{}/{}/doppio_labels/{}_labels/r   z"http://222.231.59.68/{}/{}/images/z
files.jsonZlabel_filesZimage_filesZ
mask_filesz	info.jsonbboxZ
pascal_vocclass)formatZlabel_fieldsZbbox_paramsbboxessegmentationmaskskeypointZxykeypoint_labelsZkeypoint_params	keypointskeypoint_visiblec             S   s   g | ]}|qS  r   ).0ir   r   "./doppio_dataset_client_version.py
<listcomp>D   s    z*DoppioDataset.__init__.<locals>.<listcomp>keypoint_classessegmentation_maskmaskclasses   2   
annotationimage)"dicttransforms_kwargs_inputr	   listkeysurl_image_addressurl_mask_addressospathjoinurllibZrequestZurlopenjsonloadZ	url_imageZurl_mask
json_files	info_jsonvaluestaskAZ
BboxParamsZKeypointParamsrangelenr   zipZ
classes_totalori2valid_class_num	enumerateappendZCompose	transform)selfdataset	split_dir	task_infoaugmentationZ
valid_classesZtransforms_kwargs_paramsZurl_label_addressZ	url_filesZurlZurl_files_dictZurl_temp_label_jsonZurl_infoitemZsemantic_mask_flagtZ
t_valid_classZurl_jsonZ
url_json_pathZlabel_json_fileflagannr   r   r   r   __init__   sz    







 
$

zDoppioDataset.__init__c          
   C   s.  t  | _g | jd< d | jd< x| jD ]z}|dkr<g | jd< q$|dkrPg | jd< q$|dkrxg | jd< g | jd< g | jd< q$|dkrg | jd	< q$|d
kr$g | jd
< q$W | j| }|d }tj| j|}tj	t
t|j
tjd}t|tj}t|tj}|d
 }	d
| j kr"| jd
d t| j }
x|	D ]}x|
D ]}|| krF|| }
|d
kr|
dkr|
| jd kr| j| t| jd |
  nP qF|dkr| j| |
 qF|dkrX|j\}}}tj||gtjd}tt|
jdkr.| j| t|t|
gd n(dd |
D }| j| t||d nr|d	krtj| j |d
 d d	 }tj	t
t|j
tjd}t|tj}| j| | n| j| |
 qFW q8W xB| j!D ]6}|dkr|| j!|< q|dkr| jd | j!|< n|d
kr6| jd
 | j!|< n|dkr| jd	 d }t"|}x$| j#$ D ]\}}||||k< qdW || j!|< n|dkr| jd | j!|< np|dkrdd | jd D | j!|< nJ|dkr| j%t| jd  | j!|< n$|dkrdd | jd D | j!|< qW | j&f | j!}|S )NZarear   r   r   r
   r   r   r   Zsegmentation_pathr   Z
image_name)Zdtyper   F r      r   c             S   s   g | ]}t |qS r   )nparray)r   lr   r   r   r      s    z-DoppioDataset.__getitem__.<locals>.<listcomp>r   r
   r   r   r   c             S   s   g | ]}|D ]}|qqS r   r   )r   kpelr   r   r   r      s    c             S   s   g | ]}|D ]}|qqS r   r   )r   rE   rF   r   r   r   r      s    )'r   Z
doppio_outputr-   r*   r$   r%   r&   r"   rB   Zasarray	bytearrayrequestsgetZcontentZuint8cv2ZimdecodeZIMREAD_COLORZcvtColorZ
COLOR_BGR2RGBr!   move_to_endr    r+   r4   intshapeZzerosZint32r0   rC   ZfillPolyr#   r   Z
zeros_liker2   itemsr   r5   )r6   indexr;   Z	json_fileZimg_nameZim_pathZ
image_nparrayZimgZim_rgbZannotationsZann_keysr>   Zann_keyZtemp_hw_mZmake_temp_arrayZ	mask_pathr   keyZ	copy_maskZtemp_keyvalueZ
img_tensorr   r   r   __getitem__{   s    





 

$









 zDoppioDataset.__getitem__c             C   s
   t | jS )N)r0   r*   )r6   r   r   r   __len__   s    zDoppioDataset.__len__)NNNNN)__name__
__module____qualname__r?   rV   rW   r   r   r   r   r      s   
g_r   c                s  | d   }i }x|D ]ƉdkrBfdd| D }||< qdkr|t| d d d fdd| D }||< qdkrt| d d d fd	d| D }||< qd
krt| d d d fdd| D }||< qdkrfd
d| D }t|d|< qdkrry8| d  d j\  fdd| D }||< W n   wY nX qdkrfdd| D }t|d|< qdkrƇfdd| D }||< qfdd| D }||< qW |S )Nr   Zclass_labelsc                s   g | ]}|  qS r   r   )r   s)kr   r   r      s    zcollate.<locals>.<listcomp>r   r   r   c          	      s*   g | ]"}t t|  d dqS )rA   )torch	as_tensorrB   rC   reshape)r   r[   )r\   	max_valuer   r   r      s    c                s(   g | ] }t t|  d qS )r]   )r^   r_   rB   rC   r`   )r   r[   )r\   ra   r   r   r      s    r   c                s(   g | ] }t t|  d qS )r]   )r^   r_   rB   rC   r`   )r   r[   )r\   ra   r   r   r      s    r   c                s   g | ]}|  qS r   r   )r   r[   )r\   r   r   r      s    r   c          	      s2   g | ]*}t t| tjd  qS )r]   )r^   Z
from_numpyrB   ZvstackZastypefloatr`   )r   r[   )rP   r\   rQ   r   r   r      s    r   c                s   g | ]}|  qS r   r   )r   r[   )r\   r   r   r      s    r
   c                s   g | ]}t |  qS r   )r^   r_   )r   r[   )r\   r   r   r     s    c                s   g | ]}t |  qS r   )r^   r_   )r   r[   )r\   r   r   r     s    )r!   maxr^   stackrM   )Z
batch_dataZkeys_Zreturn_dictZtmpsZtmpr   )rP   r\   ra   rQ   r   collate   sJ    









re   c             C   sL   |d kr t jdddtddg}t| |||d}tjjj|dddtd}|S )	Ni  )widthZheightT)Ztranspose_mask)r7   r8   r9   r:   rA   r   )r7   Z
batch_sizeZshuffleZnum_workersZ
collate_fn)	r.   ZResizer   r   r^   ZutilsdataZ
DataLoaderre   )r7   r8   r9   r:   Zcustom_datasetZmy_dataset_loaderr   r   r   doppio	  s    
rh   )NNNN)r(   rJ   Zalbumentationsr.   Ztorch.utils.data.datasetr   r^   Zalbumentations.pytorchr   r   ZnumpyrB   collectionsr   Zurllib.requestr'   r$   rH   r   re   rh   r   r   r   r   <module>   s    M+