jonas commited on
Commit
3daea0f
1 Parent(s): 93504a4
Files changed (1) hide show
  1. app.py +27 -11
app.py CHANGED
@@ -5,24 +5,40 @@ import os
5
 
6
  import torch
7
  import ultralytics
8
-
9
 
10
- model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5_0.65map_exp7_best.pt",
11
- force_reload=False)
 
 
 
12
  model.conf = 0.20 # NMS confidence threshold
13
 
14
  path = [['img/test-image.jpg'], ['img/test-image-2.jpg']]
15
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
16
  def show_preds_image(image_path):
17
  image = cv2.imread(image_path)
18
- results = model(image_path)
19
- results.xyxy[0] # img1 predictions (tensor)
20
- results.pandas().xyxy[0] # img1 predictions (pandas)
21
- predictions = results.pred[0]
22
- boxes = predictions[:, :4] # x1, y1, x2, y2
23
- scores = predictions[:, 4]
24
- categories = predictions[:, 5]
25
-
26
  for i, det in enumerate(results.boxes.xyxy):
27
  cv2.rectangle(
28
  image,
 
5
 
6
  import torch
7
  import ultralytics
8
+ from ultralytics import YOLO
9
 
10
+
11
+ # model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5_0.65map_exp7_best.pt",
12
+ # force_reload=False)
13
+
14
+ model = YOLO("yolov5_0.65map_exp7_best.pt")
15
  model.conf = 0.20 # NMS confidence threshold
16
 
17
  path = [['img/test-image.jpg'], ['img/test-image-2.jpg']]
18
 
19
+ # def show_preds_image(image_path):
20
+ # image = cv2.imread(image_path)
21
+ # # outputs = model(source=image_path)
22
+ # # results = outputs[0].cpu().numpy()
23
+ # results = model(image_path)
24
+ # results.xyxy[0] # img1 predictions (tensor)
25
+ # results.numpy().xyxy[0] # img1 predictions (pandas)
26
+ # predictions = results.pred[0]
27
+ # for i, det in enumerate(results.boxes.xyxy):
28
+ # cv2.rectangle(
29
+ # image,
30
+ # (int(det[0]), int(det[1])),
31
+ # (int(det[2]), int(det[3])),
32
+ # color=(0, 0, 255),
33
+ # thickness=2,
34
+ # lineType=cv2.LINE_AA
35
+ # )
36
+ # return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
37
+
38
  def show_preds_image(image_path):
39
  image = cv2.imread(image_path)
40
+ outputs = model.predict(source=image_path)
41
+ results = outputs[0].cpu().numpy()
 
 
 
 
 
 
42
  for i, det in enumerate(results.boxes.xyxy):
43
  cv2.rectangle(
44
  image,