working version

app.py

@@ -3,12 +3,12 @@ import torch
 from PIL import Image
 
 # Images
-torch.hub.download_url_to_file('https://storage.googleapis.com/kili-datasets-public/plastic-in-river/ckze0btj10ejf0lyy1imtdy7o.jpg', 'bottles1.jpg')
-torch.hub.download_url_to_file('https://storage.googleapis.com/kili-datasets-public/plastic-in-river/ckze0btj10ejd0lyyfzm85k9u.jpg', 'bottles2.jpg')
+torch.hub.download_url_to_file('https://storage.googleapis.com/kili-datasets-public/plastic-in-river/samples/ckze0btj10ejf0lyy1imtdy7o.jpg', 'bottles1.jpg')
+torch.hub.download_url_to_file('https://storage.googleapis.com/kili-datasets-public/plastic-in-river/samples/ckze0btj10ejd0lyyfzm85k9u.jpg', 'bottles2.jpg')
 
 # Model
-model = torch.hub.load_state_dict_from_url("gs://kili-datasets-public/plastic_in_river/model/best.pt", force_reload=True)  # force_reload=True to update
-# -> load the model from HF models.
+model = torch.hub.load('PierreLeveau/yolov5', 'custom', 'https://storage.googleapis.com/kili-datasets-public/plastic-in-river/model/best.pt')  # force_reload=True to update
+
 
 def yolo(im, size=640):
     g = (size / max(im.size))  # gain
@@ -22,9 +22,9 @@ def yolo(im, size=640):
 inputs = gr.inputs.Image(type='pil', label="Original Image")
 outputs = gr.outputs.Image(type="pil", label="Output Image")
 
-title = "YOLOv5"
-description = "YOLOv5 Gradio demo for object detection. Upload an image or click an example image to use."
-article = "<p style='text-align: center'>YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset, and includes simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> |<a href='https://apps.apple.com/app/id1452689527'>iOS App</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>"
+title = "YOLOv5 - Plastic in river detection"
+description = "This space demontrates a YOLOv5 model fine-tuned on a dataset of annotated photos of plastic waste in rivers. Upload an image or click an example image to use."
+article = "<p style='text-align: center'>YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset. We performed fine-tuning of models trained by Ultralytics with the help of their awesome <a href='https://github.com/ultralytics/yolov5'>code repository</a>. The data comes from a community challenge organized by [Kili](https://kili-technology.com/blog/kili-s-community-challenge-plastic-in-river-dataset), and the demo site is heavily inspired from the original [YoloV5 space](https://huggingface.co/spaces/akhaliq/YOLOv5). We will update the model during the challenge."
 
 examples = [['bottles1.jpg'], ['bottles2.jpg']]
 gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, theme="huggingface").launch(

requirements.txt

@@ -0,0 +1,32 @@
+# pip install -r requirements.txt
+
+# base ----------------------------------------
+matplotlib>=3.2.2
+numpy>=1.18.5
+opencv-python-headless
+Pillow
+PyYAML>=5.3.1
+scipy>=1.4.1
+torch>=1.7.0
+torchvision>=0.8.1
+tqdm>=4.41.0
+kili
+
+# logging -------------------------------------
+tensorboard>=2.4.1
+# wandb
+
+# plotting ------------------------------------
+seaborn>=0.11.0
+pandas
+
+# export --------------------------------------
+# coremltools>=4.1
+# onnx>=1.9.0
+# scikit-learn==0.19.2  # for coreml quantization
+
+# extras --------------------------------------
+# Cython  # for pycocotools https://github.com/cocodataset/cocoapi/issues/172
+# pycocotools>=2.0  # COCO mAP
+# albumentations>=1.0.3
+thop  # FLOPs computation