freddyaboulton HF staff commited on
Commit
385e56e
·
1 Parent(s): 8337710
Files changed (1) hide show
  1. app.py +10 -9
app.py CHANGED
@@ -4,18 +4,19 @@ import cv2
4
  import tempfile
5
  from ultralytics import YOLOv10
6
 
7
- model = YOLOv10.from_pretrained(f'jameslahm/yolov10n')
 
8
 
9
  @spaces.GPU
10
  def yolov10_inference(image, conf_threshold):
11
- width, _ = image.size
12
- import time
13
- start = time.time()
14
- results = model.predict(source=image, imgsz=width, conf=conf_threshold)
15
- end = time.time()
16
- print("time", end - start)
17
- annotated_image = results[0].plot()
18
- return annotated_image[:, :, ::-1]
19
 
20
 
21
  def app():
 
4
  import tempfile
5
  from ultralytics import YOLOv10
6
 
7
+ image_processor = RTDetrImageProcessor.from_pretrained("PekingU/rtdetr_r50vd")
8
+ model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd")
9
 
10
  @spaces.GPU
11
  def yolov10_inference(image, conf_threshold):
12
+
13
+ inputs = image_processor(images=image, return_tensors="pt")
14
+
15
+ with torch.no_grad():
16
+ outputs = model(**inputs)
17
+
18
+ results = image_processor.post_process_object_detection(outputs, target_sizes=torch.tensor([image.size[::-1]]), threshold=0.3)
19
+
20
 
21
 
22
  def app():