nielsr HF staff commited on
Commit
80244a2
1 Parent(s): 7c40d47

Load using from_pretrained

Browse files

This PR simplifies the loading of the checkpoints by leveraging the PyTorchModelHubMixin class.

Files changed (1) hide show
  1. app.py +8 -13
app.py CHANGED
@@ -2,12 +2,7 @@ import gradio as gr
2
  from ultralytics import YOLOv10
3
  import supervision as sv
4
  import spaces
5
- from huggingface_hub import hf_hub_download
6
 
7
-
8
- def download_models(model_id):
9
- hf_hub_download("kadirnar/Yolov10", filename=f"{model_id}", local_dir=f"./")
10
- return f"./{model_id}"
11
 
12
  box_annotator = sv.BoxAnnotator()
13
  category_dict = {
@@ -33,7 +28,7 @@ category_dict = {
33
  @spaces.GPU(duration=200)
34
  def yolov10_inference(image, model_id, image_size, conf_threshold, iou_threshold):
35
  model_path = download_models(model_id)
36
- model = YOLOv10(model_path)
37
  results = model(source=image, imgsz=image_size, iou=iou_threshold, conf=conf_threshold, verbose=False)[0]
38
  detections = sv.Detections.from_ultralytics(results)
39
 
@@ -54,14 +49,14 @@ def app():
54
  model_id = gr.Dropdown(
55
  label="Model",
56
  choices=[
57
- "yolov10n.pt",
58
- "yolov10s.pt",
59
- "yolov10m.pt",
60
- "yolov10b.pt",
61
- "yolov10l.pt",
62
- "yolov10x.pt",
63
  ],
64
- value="yolov10m.pt",
65
  )
66
  image_size = gr.Slider(
67
  label="Image Size",
 
2
  from ultralytics import YOLOv10
3
  import supervision as sv
4
  import spaces
 
5
 
 
 
 
 
6
 
7
  box_annotator = sv.BoxAnnotator()
8
  category_dict = {
 
28
  @spaces.GPU(duration=200)
29
  def yolov10_inference(image, model_id, image_size, conf_threshold, iou_threshold):
30
  model_path = download_models(model_id)
31
+ model = YOLOv10.from_pretrained(f"jameslahm/{model_id}")
32
  results = model(source=image, imgsz=image_size, iou=iou_threshold, conf=conf_threshold, verbose=False)[0]
33
  detections = sv.Detections.from_ultralytics(results)
34
 
 
49
  model_id = gr.Dropdown(
50
  label="Model",
51
  choices=[
52
+ "yolov10n",
53
+ "yolov10s",
54
+ "yolov10m",
55
+ "yolov10b",
56
+ "yolov10l",
57
+ "yolov10x",
58
  ],
59
+ value="yolov10m",
60
  )
61
  image_size = gr.Slider(
62
  label="Image Size",