FPh#ddlZddlZddlmZddlZddlZddlmZddl m Z m Z ddl m Z mZmZmZmZmZmZmZddlmZmZddlmZmZddlmZd d lmZd d lm Z Gd d e jBZ"GddZ#dZ$ddZ%ddZ&dZ'ddZ(y)N)Path)Image) dataloader distributed)LOADERS LoadImagesLoadPilAndNumpyLoadScreenshots LoadStreams LoadTensor SourceTypes autocast_list) IMG_FORMATS VID_FORMATS)RANKcolorstr) check_file) YOLODataset) PIN_MEMORYc4eZdZdZfdZdZdZdZxZS)InfiniteDataLoaderzV Dataloader that reuses workers. Uses same syntax as vanilla DataLoader. ct||i|tj|dt |j t||_y)zFDataloader that infinitely recycles workers, inherits from DataLoader. batch_samplerN)super__init__object __setattr___RepeatSamplerr__iter__iterator)selfargskwargs __class__s aC:\Users\daisl\Desktop\realtime-object-detection\venv\Lib\site-packages\ultralytics/data/build.pyrzInfiniteDataLoader.__init__sA $)&)4.ASAS2TU(* c@t|jjS)z2Returns the length of the batch sampler's sampler.)lenrsamplerr"s r&__len__zInfiniteDataLoader.__len__#s4%%--..r'c#lKtt|D]}t|jyw)z,Creates a sampler that repeats indefinitely.N)ranger)nextr!)r"_s r&r zInfiniteDataLoader.__iter__'s(s4y!At}}% %"s24c.|j|_y)zt Reset iterator. This is useful when we want to modify settings of dataset while training. N) _get_iteratorr!r+s r&resetzInfiniteDataLoader.reset,s **, r') __name__ __module__ __qualname____doc__rr,r r3 __classcell__)r%s@r&rrs + /& -r'rceZdZdZdZdZy)rzl Sampler that repeats forever. Args: sampler (Dataset.sampler): The sampler to repeat. c||_y)z@Initializes an object that repeats a given sampler indefinitely.N)r*)r"r*s r&rz_RepeatSampler.__init__=s  r'c#LK t|jEd{7w)z4Iterates over the 'sampler' and yields its contents.N)iterr*r+s r&r z_RepeatSampler.__iter__As#DLL) ) ) )s $"$N)r4r5r6r7rr r'r&rr5s*r'rctjdz}tjj |tj|y)z\Set dataloader worker seed https://pytorch.org/docs/stable/notes/randomness.html#dataloader.lN)torch initial_seednprandomseed) worker_id worker_seeds r& seed_workerrFGs1$$&0KIINN; KK r'c tdid|d|jd|d|dk(d|d|jxs|d|jxsd d |jxsd d t |d |dk(rdnddt |dd|jdk(d|jdk(d|jd|d|dk(r|jSdS)zBuild YOLO Dataset.img_pathimgsz batch_sizeaugmenttrainhyprectcacheN single_clsFstridepadgg?prefixz: use_segmentssegment use_keypointsposeclassesdatafractiong?r=) rrIrNrOrPintrtaskrXrZ)cfgrHbatchrYmoderNrQs r&build_yolo_datasetr`Ns  ; ;ii ; ;  ;   ; XX   ;ii4 ;>>*U ;6{ ;7?C ;4& $ ;XX* ;hh&( ;  ; ; "&! ;; 7:! ;;r'c t|t|}tjj }tt j t|dz|dkDr|nd|g}|dk(rdntj||}tj}|jdtzt|||xr|du||tt|ddt | S) zJReturn an InfiniteDataLoader or DataLoader for training or validation set.rrN)shufflelUU*UU* collate_fn) datasetrJrc num_workersr* pin_memoryrdworker_init_fn generator)minr)r?cuda device_countos cpu_countmaxrDistributedSampler Generator manual_seedrrrgetattrrF) rer^workersrcrankndnwr*ris r&build_dataloaderrxcs s7| $E  "B bllnB *UQYEAw O PBbjdk&D&DWV]&^G!I -45 g).&-&A'T/*,&-)3)0,)M-8(1 33r'cd\}}}}}t|tttfrt|}t|jddt t zv}|jjd}|jxs|jdxs|xr| }|jdk(}|r|rt|}nt|trd}nyt|ttfrt|}d}nUt|t j t"j$frd}n(t|t&j(rd}n t+d||||||fS) z7Check source type and return corresponding flag values.)FFFFFrN)zhttps://zhttp://zrtsp://zrtmp://ztcp://z.streamsscreenTzZUnsupported image type. For supported types see https://docs.ultralytics.com/modes/predict) isinstancestrr[rsuffixrrlower startswith isnumericendswithrrlisttuplerrrAndarrayr?Tensor TypeError)sourcewebcam screenshotfrom_img in_memorytensoris_fileis_urls r& check_sourcervs&6W3FJ)V&3T*+Vv,%%ab)kK.GH**+bc!!#^vz'B^vG]V]R]\\^x/ g'F FG $ FT5M *v& FU[["**5 6 FELL )tuu 6:xF BBr'c"t|\}}}}}}|r |jn t||||} |r t|} nE|r|} n@|rt ||||} n.|rt ||} n|rt ||} nt|||} t| d| | S)a1 Loads an inference source for object detection and applies necessary transformations. Args: source (str, Path, Tensor, PIL.Image, np.ndarray): The input source for inference. imgsz (int, optional): The size of the image for inference. Default is 640. vid_stride (int, optional): The frame interval for video sources. Default is 1. buffer (bool, optional): Determined whether stream frames will be buffered. Default is False. Returns: dataset (Dataset): A dataset object for the specified input source. )rI vid_stridebuffer)rI)rIr source_type) rrr r r r r rsetattr) rrIrrrrrrrrres r&load_inference_sourcers?K6>R;FFJ)V(1&$${6:W_ag7hKV$  fEjQWX !&6 !&6V5ZH G]K0 Nr')rLF )Trb)NirF))rmrBpathlibrnumpyrAr?PILrtorch.utils.datarrultralytics.data.loadersrrr r r r r rultralytics.data.utilsrrultralytics.utilsrrultralytics.utils.checksrrerutilsr DataLoaderrrrFr`rxrrr=r'r&rss  4BBB;,/ -..->**$;*3&C4!r'