Update inference.py

inference.py

@@ -1,17 +1,13 @@
-import torch
 from roboflow import Roboflow
-from pathlib import Path
+from PIL import Image
+import numpy as np
 
-def run_detection(image_path):
-    rf = Roboflow(api_key="rSfidObyvZ1BYV7yOQrN")
-    project = rf.workspace("naveen-kumar-9emxl").project("my-first-project-icheu")
-    dataset = project.version(1).download("yolov5")
+# Load Roboflow model
+rf = Roboflow(api_key="rSfidObyvZ1BYV7yOQrN")
+project = rf.workspace("naveen-kumar-9emxl").project("my-first-project-icheu")
+model = project.version(1).model
 
-    model_path = "runs/train/exp/weights/best.pt"
-    if not Path(model_path).exists():
-        print("Trained model not found.")
-        return
-
-    model = torch.hub.load('yolov5', 'custom', path=model_path, source='local')
-    results = model(image_path)
-    results.show()
+def predict(image):
+    image = np.array(image)
+    prediction = model.predict(image).plot()
+    return Image.fromarray(prediction)