Add examples and model

app.py

@@ -17,26 +17,24 @@ from utils.general import check_img_size, check_requirements, check_imshow, non_
 from utils.plots import plot_one_box
 from utils.torch_utils import select_device, load_classifier, time_synchronized, TracedModel
 from PIL import Image
- 
+
 from huggingface_hub import hf_hub_download
 
 
 def load_model(model_name):
-    model_path = hf_hub_download(repo_id=f"Yolov7/{model_name}", filename=f"{model_name}.pt")
+    model_path = hf_hub_download(repo_id="mw00/yolov7-lego", filename=f"{model_name}.pt")
 
     return model_path
 
 
 model_names = [
-    "yolov7",
-    "yolov7-e6e",
-    "yolov7-e6",
+    "best",
 ]
 
 models = {model_name: load_model(model_name) for model_name in model_names}
 
 
-def detect(img,model):
+def detect(img, model):
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', nargs='+', type=str, default=models[model], help='model.pt path(s)')
     parser.add_argument('--source', type=str, default='Inference/', help='source')  # file/folder, 0 for webcam
@@ -159,7 +157,7 @@ def detect(img,model):
                         plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
 
             # Print time (inference + NMS)
-            #print(f'{s}Done. ({t2 - t1:.3f}s)')
+            # print(f'{s}Done. ({t2 - t1:.3f}s)')
 
             # Stream results
             if view_img:
@@ -187,11 +185,13 @@ def detect(img,model):
 
     if save_txt or save_img:
         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
-        #print(f"Results saved to {save_dir}{s}")
+        # print(f"Results saved to {save_dir}{s}")
 
     print(f'Done. ({time.time() - t0:.3f}s)')
 
-    return Image.fromarray(im0[:,:,::-1])
+    return Image.fromarray(im0[:, :, ::-1])
+
 
-   
-gr.Interface(detect,[gr.Image(type="pil"),gr.Dropdown(choices=model_names)], gr.Image(type="pil"),title="Yolov7",examples=[["horses.jpeg", "yolov7"]],description="demo for <a href='https://github.com/WongKinYiu/yolov7' style='text-decoration: underline' target='_blank'>WongKinYiu/yolov7</a> Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors").launch()
+gr.Interface(detect, [gr.Image(type="pil"), gr.Dropdown(choices=model_names)], gr.Image(type="pil"), title="Yolov7",
+             examples=[["lego-1.jpg", "yolov7", "lego-2.jpg", "yolov7"]],
+             description="demo for <a href='https://github.com/WongKinYiu/yolov7' style='text-decoration: underline' target='_blank'>WongKinYiu/yolov7</a> Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors").launch()