Upload New File

iot/InspBot/yolov5.py

+import io
+import torch
+from PIL import Image
+
+ddWeights = './weights/dd.pt'
+gsWeights = './weights/gs.pt'
+
+
+def setModel(modelType):
+    global model
+    weight = ddWeights
+    if modelType == 'gs':
+        weight = gsWeights
+    model = torch.hub.load("./yolov5", 'custom', path=weight, source='local')
+    model.conf = 0.6
+
+
+def yolov5(img):
+    results = model(img)
+    detected_classes = []
+    confidence = 0.0
+    names = results.names
+    if results.pred is not None:
+        pred = results.pred[0]
+        if pred is not None:
+            for c in pred[:, -1].unique():
+                n = (pred[:, -1] == c).sum()
+                confidence = pred[0][4].item()
+                detected_classes.append(f"{names[int(c)]}{'s' * (n > 1)}")
+
+    rendered_imgs = results.render()
+    converted_img = Image.fromarray(rendered_imgs[0]).convert("RGB")
+    return detected_classes, confidence, converted_img