Coverage for collectra/tasks/machine_learning/orienters.py: 0.00%

1from __future__ import annotations 

2 

3from pathlib import Path 

4from typing import TYPE_CHECKING 

5 

6from collectra.types.images import Image, ImageCrop, Orientation 

7 

8from ..base import Task 

9 

10if TYPE_CHECKING: 

11 from torchvision.models import EfficientNet 

12 

13__all__ = ["ImageOrienter"] 

14 

15 

16class ImageOrienter(Task): 

17 

18 model: str | Path | EfficientNet 

19 

20 def __init__(self, name, model: str | Path | EfficientNet, **kwargs): 

21 super().__init__(name, model=model, **kwargs) 

22 self.transforms = None 

23 

24 def _init_model(self): 

25 """Initialisation code here is emulated from https://github.com/duartebarbosadev/deep-image-orientation-detection""" 

26 import torch 

27 import torchvision.transforms as transforms 

28 from torchvision.models import EfficientNet 

29 

30 image_size = 384 

31 self.transforms = transforms.Compose( 

32 [ 

33 transforms.Resize((image_size + 32, image_size + 32)), 

34 transforms.CenterCrop(image_size), 

35 transforms.ToTensor(), 

36 transforms.Normalize( 

37 mean=[0.485, 0.456, 0.406], 

38 std=[0.229, 0.224, 0.225], 

39 ), 

40 ] 

41 ) 

42 

43 device = torch.device( 

44 "mps" 

45 if torch.backends.mps.is_available() 

46 else "cuda" if torch.cuda.is_available() else "cpu" 

47 ) 

48 

49 if isinstance(self.model, EfficientNet): 

50 self.model.to(device) 

51 self.device = device 

52 return 

53 

54 import torchvision.models as models 

55 

56 model = models.efficientnet_v2_s(weights=None) 

57 

58 num_ftrs = model.classifier[1].in_features 

59 

60 model.classifier = torch.nn.Sequential( 

61 torch.nn.Dropout(p=0.2, inplace=True), 

62 torch.nn.Linear(num_ftrs, 4), 

63 ) 

64 

65 state_dict = torch.load(self.model, map_location=device) 

66 model.load_state_dict(state_dict) 

67 model.to(device) 

68 model.eval() 

69 

70 self.model = model 

71 self.device = device 

72 

73 def _processed_image(self, img: Image): 

74 from PIL import Image as PILImage 

75 

76 image = img.pil() 

77 

78 if image.mode in ("RGB", "L"): 

79 return image.convert("RGB") 

80 

81 converted_image = image.convert("RGBA") 

82 background = PILImage.new("RGBA", converted_image.size, (255, 255, 255)) 

83 background.paste(converted_image, mask=converted_image) 

84 return background 

85 

86 def run(self, *args: Image) -> Image: 

87 import torch 

88 from torchvision.models import EfficientNet 

89 

90 if len(args) != 1: 

91 raise ValueError("This task only supports a single Image input.") 

92 self._init_model() 

93 img = args[0] 

94 image = self._processed_image(img) 

95 input_tensor = self.transforms(image).unsqueeze(0).to(self.device) 

96 

97 if not isinstance(self.model, EfficientNet): 

98 raise ValueError("Model must be an instance of EfficientNet.") 

99 

100 with torch.no_grad(): 

101 output = self.model(input_tensor) 

102 _, predicted_idx = torch.max(output, 1) 

103 predicted_class = predicted_idx.item() 

104 orientation = Orientation(predicted_class) 

105 

106 name = self.get_output_name() 

107 

108 if isinstance(img, ImageCrop): 

109 new_image = ImageCrop( 

110 name=name, 

111 data=img.data, 

112 x_center=float(img.x_center), 

113 y_center=float(img.y_center), 

114 width_relative=float(img.width_relative), 

115 height_relative=float(img.height_relative), 

116 orientation=orientation, 

117 ) 

118 return new_image 

119 

120 new_image = Image( 

121 name=name, 

122 data=img.data, 

123 orientation=orientation, 

124 ) 

125 

126 return new_image