Features

Requirements

Usage

fromtta_pytorchimportTYPES,Chain,Compose,Flip,HFlip,Merger,Rescale,Resize,VFliptta_trans=Compose(    [Rescale(scales=[0.5],image_mode="bilinear",image_align_corners=False,mask_mode="bilinear",mask_align_corners=False,        ),Resize(sizes=[128],image_mode="bilinear",image_align_corners=False,mask_mode="bilinear",mask_align_corners=False,        ),Flip(),HFlip(),# VFlip(),    ],verbose=True,)image=torch.randn(3,1,50,50,dtype=torch.float32)

tta_results=Merger()fortransintta_trans:trans:Chainaug_image=trans.do_image(image)undo_image=trans.undo_image(aug_image)tta_results.append(undo_image)seg_results=tta_results.result

tta_results=Merger()fortransintta_trans:trans:Chainaug_images:List[torch.Tensor]=trans.do_all(inputs=[image],input_types=[TYPES.IMAGE])undo_images=trans.undo_all(outputs=aug_images,output_types=[TYPES.MASK])tta_results.append(undo_images[0])seg_results=tta_results.result

tta_seg_merger=Merger(mode="mean")tta_cls_merger=Merger(mode="mean")tta_seg_merger.reset()tta_cls_merger.reset()fortranintta_trans:tran:Chainaug_tensor=tran.do_image(image)# simulate real datamask=aug_tensorlabel=torch.randn(3,1000,dtype=torch.float32)# for segmentation, [B,K,H,W]undo_mask=tran.undo_image(mask)tta_seg_merger.append(undo_mask)# for classification, [B,K]undo_label=tran.undo_label(label)tta_cls_merger.append(undo_label)seg_results=tta_seg_merger.resultseg_mask=seg_results.argmax(dim=1)# [B,H,W]cls_results=tta_cls_merger.resultcls_score,cls_index=cls_results.max(dim=1)# [B], [B]

tta_seg_results= []tta_cls_results= []fortranintta_trans:tran:Chainaug_tensor=tran.do_image(image)# simulate real datamask=aug_tensorlabel=torch.randn(3,1000,dtype=torch.float32)# for segmentation, [B,K,H,W]undo_mask=tran.undo_image(mask)tta_seg_results.append(undo_mask)# for classification, [B,K]undo_label=tran.undo_label(label)tta_cls_results.append(undo_label)seg_results=sum(tta_seg_results)/len(tta_seg_results)seg_mask=seg_results.argmax(dim=1)# [B,H,W]cls_results=sum(tta_cls_results)/len(tta_cls_results)cls_score,cls_index=cls_results.max(dim=1)# [B], [B]

@tta_trans.decorate(input_infos={"image":TYPES.IMAGE},output_infos={"mask":TYPES.MASK,"label":TYPES.LABEL},merge_mode="mean",)defdo_something(image=None):label=torch.randn(3,1000,dtype=torch.float32)return {"mask":image,"label":label}tta_results=do_something(image=image)

Cite

@online{tta.pytorch,author="lartpang",title="{Test-Time Augmentation library for Pytorch}",url="https://github.com/lartpang/tta.pytorch",note="(Dec 20, 2023)",}