Update handler.py

handler.py

@@ -21,8 +21,8 @@ class EndpointHandler():
     def __init__(self, path=""):
         # Preload all the elements you are going to need at inference.
         self.dinov2_vits14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
-        device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
-        self.dinov2_vits14.to(device)
+        self.device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
+        self.dinov2_vits14.to(self.device)
         print('Successfully load dinov2_vits14 model')
         
         self.yolov8_model = YOLO(os.path.join(path, 'yolov8_2023-07-19_yolov8m.pt'))
@@ -40,6 +40,16 @@ class EndpointHandler():
 
         with open(os.path.join(path, 'labels.txt'), 'r') as f:
             self.labels = f.read().split(',') # loggerhead,green,leatherback...
+
+        self.name_en2vi = {
+            "loggerhead": "Quản đồng",
+            "green": "Vích",
+            "leatherback": "Rùa da",
+            "hawksbill": "Đồi mồi",
+            "kemp_ridley": "Vích Kemp",
+            "olive_ridley": "Đồi mồi dứa",
+            "flatback": "Rùa lưng phẳng"
+        }
     
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -73,16 +83,16 @@ class EndpointHandler():
         else:
             y1 = max(y1 - offset, 0)
             y2 = min(y2 + offset, H)
-        new_image = img[y1:y2, x1:x2]
+        cropped = img[y1:y2, x1:x2]
 
         new_image = self.transform_image(Image.fromarray(cropped))[:3].unsqueeze(0)
-        embedding = self.dinov2_vits14(new_image.to(device))
+        embedding = self.dinov2_vits14(new_image.to(self.device))
         prediction = self.linear_model(embedding)
         percentage = nn.Softmax(dim=1)(prediction).detach().numpy().round(2)[0].tolist()
         result = {}
         
         for i in range(len(self.labels)):
-            result[name_en2vi[self.labels[i]]] = percentage[i]
+            result[self.name_en2vi[self.labels[i]]] = percentage[i]
 
         # Return the annotated original image with the square cropped and result dict
-        return annotated.tolist(), result
+        return annotated.tolist(), result