refactor final

Author: aibekakhmetkazy

Dockerfile

@@ -31,9 +31,6 @@ RUN pip3 install -r requirements.txt
 # Set the working directory
 WORKDIR /usr/src/app
 
-# Run the tests to ensure everything is working correctly
-#RUN make test
-
 # Set the entrypoint to run the pipeline via Makefile.run
 ENTRYPOINT ["make", "-f", "Makefile.run"]
 

main.py

@@ -7,39 +7,39 @@
 from nets import MobileNetV2
 from assets.meta import IMAGENET_CATEGORIES
 
+
+# Assume the image is already resized; no need for transforms.Resize
+def preprocess_image(resized_image_path):
+    transform = transforms.Compose([
+        transforms.ToTensor(),  # Convert the image to a PyTorch tensor
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],  # Normalize using the mean and std of ImageNet
+            std=[0.229, 0.224, 0.225]
+        ),
+    ])
+
+    image = Image.open(resized_image_path)
+    image = transform(image).unsqueeze(0)  # Add a batch dimension
+    return image, image.shape, type(image)
+
+def inference(model, resized_image_path, output_file_path):
+    model.eval()
+
+    input_image, image_shape, image_type = preprocess_image(resized_image_path)
+    with torch.no_grad():
+        output = model(input_image)
+
+    _, predicted_class = output.max(1)
+    predicted_label = IMAGENET_CATEGORIES[predicted_class.item()]
+    # Write output in text file
+    with open(output_file_path, "w") as f:
+        f.write(predicted_label)
+    return predicted_label
+
 def main():
     model = MobileNetV2()
     model.load_state_dict(torch.load("./model/weights/mobilenetv2.pt", weights_only=True))  # weights ported from torchvision
     model.float()  # converting weights to float32
-
-    # Assume the image is already resized; no need for transforms.Resize
-    def preprocess_image(resized_image_path):
-        transform = transforms.Compose([
-            transforms.ToTensor(),  # Convert the image to a PyTorch tensor
-            transforms.Normalize(
-                mean=[0.485, 0.456, 0.406],  # Normalize using the mean and std of ImageNet
-                std=[0.229, 0.224, 0.225]
-            ),
-        ])
-
-        image = Image.open(resized_image_path)
-        image = transform(image).unsqueeze(0)  # Add a batch dimension
-        return image, image.shape, type(image)
-
-    def inference(model, resized_image_path, output_file_path):
-        model.eval()
-
-        input_image, image_shape, image_type = preprocess_image(resized_image_path)
-        with torch.no_grad():
-            output = model(input_image)
-
-        _, predicted_class = output.max(1)
-        predicted_label = IMAGENET_CATEGORIES[predicted_class.item()]
-        # Write output in text file
-        with open(output_file_path, "w") as f:
-            f.write(predicted_label)
-        return predicted_label
-
     # Input directory
     input_dir = "input"
     resized_image_path = os.path.join(input_dir, "resized_image.jpg")

test_main.py

@@ -1,39 +1,47 @@
 import unittest
 import sys
+import os
 sys.path.append("/usr/src/app/model")
 from unittest import mock
 from unittest.mock import patch
-from main import main
+import main
 from nets import MobileNetV2
 import torch
 
 class TestMainFunction(unittest.TestCase):
 
+    def __init__(self):
+        self.input_dir = "input_raw"
+        self.resized_image_path = os.path.join(self.input_dir, "image.jpg")
+        self.output_dir = "output_raw"
+        os.makedirs(self.output_dir, exist_ok=True)  # Create output_raw if it doesn't exist
+        self.output_file_path = os.path.join(self.output_dir, "output_prediction.txt")
+
     def test_missing_input_image(self):
         # Test with a missing input image
         with self.assertRaises(FileNotFoundError):
-            main().preprocess_image('/nonexistent/image.jpg')
+            main.preprocess_image('/nonexistent/image.jpg')
 
     def test_tensor_size(self):
-        image, img_shape, img_type = main().preprocess_image('/input/resized_image.jpg')
+        image, img_shape, img_type = main.preprocess_image(self.resized_image_path)
         print(img_shape, img_type)
-        self.assertNotEquals(img_shape, (1,2,3), 'Not Equal Shape')
+        self.assertNotEqual(img_shape, (1,2,3), 'Equal Shape')
 
     def test_image_type(self):
-        image, img_shape, img_type = main().preprocess_image('/input/resized_image.jpg')
+        image, img_shape, img_type = main.preprocess_image(self.resized_image_path)
         print(img_shape, img_type)
-        self.assertNotEquals(img_type, int, 'Not Equal Image Type')
+        self.assertNotEqual(img_type, int, 'Equal Image Type')
 
     def test_image_class(self):
         model = MobileNetV2()
-        model.load_state_dict(torch.load("./model/weights/mobilenetv2.pt", weights_only=True))  # weights ported from torchvision
+        model.load_state_dict(torch.load("/Users/aibekakhmetkazy/PycharmProjects/fse4ai_team_2/mobilenetv2-pytorch/weights/mobilenetv2.pt", weights_only=True))  # weights ported from torchvision
         model.float()
 
-        predicted_label = main().inference(model, '/input/resized_image.jpg',
-                         '/output_raw/output_prediction.txt')
+        predicted_label = main.inference(model, self.resized_image_path,
+                         self.output_file_path)
 
         print('Label in output:',predicted_label)
-        self.assertNotEquals(predicted_label, 'Aar', 'Wrong classification')
+        self.assertNotEqual(predicted_label, 'House', 'Correct classification')
 
 if __name__ == '__main__':
     unittest.main()