Coverage for src/bob/pad/face/database/database.py: 97%

1import os 

2 

3from functools import partial 

4 

5from sklearn.preprocessing import FunctionTransformer 

6 

7from bob.bio.base.utils.annotations import read_annotation_file 

8from bob.bio.video import VideoAsArray 

9from bob.pipelines import DelayedSample 

10 

11 

12def get_no_transform(x): 

13 return None 

14 

15 

16def delayed_video_load( 

17 samples, 

18 original_directory, 

19 annotation_directory=None, 

20 selection_style=None, 

21 max_number_of_frames=None, 

22 step_size=None, 

23 get_transform=None, 

24 keep_extension_for_annotation=False, 

25): 

26 get_transform = get_transform or get_no_transform 

27 

28 original_directory = original_directory or "" 

29 annotation_directory = annotation_directory or "" 

30 results = [] 

31 for sample in samples: 

32 video_path = os.path.join(original_directory, sample.filename) 

33 data = partial( 

34 VideoAsArray, 

35 path=video_path, 

36 selection_style=selection_style, 

37 max_number_of_frames=max_number_of_frames, 

38 step_size=step_size, 

39 transform=get_transform(sample), 

40 ) 

41 delayed_attributes = {"annotations": lambda: None} 

42 if annotation_directory: 

43 path = sample.filename 

44 if not keep_extension_for_annotation: 

45 path = os.path.splitext(sample.filename)[0] 

46 

47 # Check if the annotation directory is a local directory or a tar.gz 

48 if os.path.isdir(annotation_directory): 

49 file_name = os.path.join(annotation_directory, path + ".json") 

50 else: 

51 file_name = f"{annotation_directory}:{path}.json" 

52 

53 delayed_annotations = partial( 

54 read_annotation_file, 

55 file_name=file_name, 

56 annotation_type="json", 

57 ) 

58 delayed_attributes = {"annotations": delayed_annotations} 

59 if sample.attack_type == "": 

60 sample.attack_type = None 

61 sample.is_bonafide = sample.attack_type is None 

62 if not hasattr(sample, "key"): 

63 sample.key = sample.filename 

64 

65 results.append( 

66 DelayedSample( 

67 data, 

68 parent=sample, 

69 delayed_attributes=delayed_attributes, 

70 ) 

71 ) 

72 return results 

73 

74 

75def VideoPadSample( 

76 original_directory, 

77 annotation_directory=None, 

78 selection_style=None, 

79 max_number_of_frames=None, 

80 step_size=None, 

81 get_transform=None, 

82 keep_extension_for_annotation=False, 

83): 

84 return FunctionTransformer( 

85 delayed_video_load, 

86 validate=False, 

87 kw_args=dict( 

88 original_directory=original_directory, 

89 annotation_directory=annotation_directory, 

90 selection_style=selection_style, 

91 max_number_of_frames=max_number_of_frames, 

92 step_size=step_size, 

93 get_transform=get_transform, 

94 keep_extension_for_annotation=keep_extension_for_annotation, 

95 ), 

96 )