navoditamathur commited on
Commit
c58fa6f
·
1 Parent(s): 8f1ccd0
Files changed (5) hide show
  1. .DS_Store +0 -0
  2. models/.DS_Store +0 -0
  3. models/best.pt +3 -0
  4. requirements.txt +3 -1
  5. tasks/image.py +28 -24
.DS_Store ADDED
Binary file (6.15 kB). View file
 
models/.DS_Store ADDED
Binary file (6.15 kB). View file
 
models/best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f901a7a25694ac87ad2ed3a047f5118448c23a4583a8326861daa2104c26218d
3
+ size 6224483
requirements.txt CHANGED
@@ -7,4 +7,6 @@ pydantic>=1.10.0
7
  python-dotenv>=1.0.0
8
  gradio>=4.0.0
9
  requests>=2.31.0
10
- librosa==0.10.2.post1
 
 
 
7
  python-dotenv>=1.0.0
8
  gradio>=4.0.0
9
  requests>=2.31.0
10
+ librosa==0.10.2.post1
11
+ torch
12
+ ultralytics
tasks/image.py CHANGED
@@ -5,7 +5,7 @@ import numpy as np
5
  from sklearn.metrics import accuracy_score, precision_score, recall_score
6
  import random
7
  import os
8
-
9
  from .utils.evaluation import ImageEvaluationRequest
10
  from .utils.emissions import tracker, clean_emissions_data, get_space_info
11
 
@@ -100,36 +100,40 @@ async def evaluate_image(request: ImageEvaluationRequest):
100
  # Update the code below to replace the random baseline with your model inference
101
  #--------------------------------------------------------------------------------------------
102
 
 
 
 
103
  predictions = []
104
  true_labels = []
105
- pred_boxes = []
106
- true_boxes_list = [] # List of lists, each inner list contains boxes for one image
107
-
108
  for example in test_dataset:
109
- # Parse true annotation (YOLO format: class_id x_center y_center width height)
 
110
  annotation = example.get("annotations", "").strip()
 
111
  has_smoke = len(annotation) > 0
112
- true_labels.append(int(has_smoke))
113
-
114
- # Make random classification prediction
115
- pred_has_smoke = random.random() > 0.5
116
- predictions.append(int(pred_has_smoke))
117
-
118
- # If there's a true box, parse it and make random box prediction
119
  if has_smoke:
120
- # Parse all true boxes from the annotation
121
  image_true_boxes = parse_boxes(annotation)
122
- true_boxes_list.append(image_true_boxes)
123
-
124
- # For baseline, make one random box prediction per image
125
- # In a real model, you might want to predict multiple boxes
126
- random_box = [
127
- random.random(), # x_center
128
- random.random(), # y_center
129
- random.random() * 0.5, # width (max 0.5)
130
- random.random() * 0.5 # height (max 0.5)
131
- ]
132
- pred_boxes.append(random_box)
 
 
 
 
 
 
133
 
134
  #--------------------------------------------------------------------------------------------
135
  # YOUR MODEL INFERENCE STOPS HERE
 
5
  from sklearn.metrics import accuracy_score, precision_score, recall_score
6
  import random
7
  import os
8
+ from ultralytics import YOLO
9
  from .utils.evaluation import ImageEvaluationRequest
10
  from .utils.emissions import tracker, clean_emissions_data, get_space_info
11
 
 
100
  # Update the code below to replace the random baseline with your model inference
101
  #--------------------------------------------------------------------------------------------
102
 
103
+ model = YOLO("../models/best.pt")
104
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
105
+ model = model.to(device)
106
  predictions = []
107
  true_labels = []
108
+ pred_boxes = []
109
+ true_boxes_list = []
 
110
  for example in test_dataset:
111
+ # Extract image and annotations
112
+ image = example["image"]
113
  annotation = example.get("annotations", "").strip()
114
+
115
  has_smoke = len(annotation) > 0
116
+ true_labels.append(1 if has_smoke else 0)
117
+
 
 
 
 
 
118
  if has_smoke:
 
119
  image_true_boxes = parse_boxes(annotation)
120
+ if image_true_boxes:
121
+ true_boxes_list.append(image_true_boxes)
122
+ else:
123
+ true_boxes_list.append([])
124
+ else:
125
+ true_boxes_list.append([])
126
+
127
+ results = model .predict(image, verbose=False) #
128
+
129
+ if len(results[0].boxes):
130
+ pred_box = results[0].boxes.xywhn[0].cpu().numpy().tolist()
131
+ predictions.append(1)
132
+ pred_boxes.append(pred_box)
133
+ else:
134
+ predictions.append(0)
135
+ pred_boxes.append([])
136
+
137
 
138
  #--------------------------------------------------------------------------------------------
139
  # YOUR MODEL INFERENCE STOPS HERE