From 29906575c39bfc2f4dfe1fd11a4d68232b01c460 Mon Sep 17 00:00:00 2001
From: KOSASIH <kosasihg88@gmail.com>
Date: Fri, 5 Jul 2024 16:44:34 +0700
Subject: [PATCH] Create ResNet.py

---
 .ai/computer_vision/ResNet.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 .ai/computer_vision/ResNet.py

diff --git a/.ai/computer_vision/ResNet.py b/.ai/computer_vision/ResNet.py
new file mode 100644
index 0000000..06faeda
--- /dev/null
+++ b/.ai/computer_vision/ResNet.py
@@ -0,0 +1,14 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+class ResNet:
+    def __init__(self, model_path):
+        self.model = models.resnet50(pretrained=True)
+        self.model.load_state_dict(torch.load(model_path))
+
+    def predict(self, image):
+        image = image.unsqueeze(0)
+        outputs = self.model(image)
+        _, predicted = torch.max(outputs, 1)
+        return predicted.item()