From 759b403042ad58b3ab4742a9819684f78b04c716 Mon Sep 17 00:00:00 2001
From: Luca <luca.medeiros@nuvi-labs.com>
Date: Tue, 5 Jan 2021 19:11:26 +0900
Subject: [PATCH] Fix main and add documentation

---
 README.md | 23 +++++++++++++++++++++++
 main.py   |  4 ++--
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 7477da7..bd71488 100644
--- a/README.md
+++ b/README.md
@@ -1,3 +1,26 @@
 # AGC - Nuvi RecycleNet
 Effective trash detection model for AI Grand Challenge 2020 track 3 round 1.  
 Nuvilabs Solution.
+The classes to be detected are:
+
+ - combustion 
+ - paper 
+ - steel 
+ - glass 
+ - plastic 
+ - plasticbag 
+ - styrofoam 
+ - food
+
+## Usage
+Download the pretrained model
+
+    wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1LnCw4L8RnM3qM96czxTTsOHn3xAKEpya' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id=1LnCw4L8RnM3qM96czxTTsOHn3xAKEpya" -O ./model/model_checkpoint.pth && rm -rf /tmp/cookies.txt
+
+
+Run inference
+
+    python main.py --img_path ./sample_img.jpg
+Will return a dictionary with the following format:
+{'Annotations': [{'Label': 'steel', 'Bbox': [564, 383, 695, 443], 'Confidence': 0.9856872}]}
+The bbox format follows: [xmin, ymin, xmax, ymax]
diff --git a/main.py b/main.py
index ae16f7f..c268e3a 100755
--- a/main.py
+++ b/main.py
@@ -8,7 +8,7 @@
 import argparse
 import os.path as osp
 import glob
-from drs_infer import Nuvi_RecycleNet
+from inference import Nuvi_RecycleNet
 
 
 parser = argparse.ArgumentParser(
@@ -19,7 +19,7 @@
                     help='Model checkpoint file path.')
 parser.add_argument('--img_path', default='./sample_img.jpg', type=str,
                     help='Path of image or images.')
-parser.add_argument('--threshold', default=0.5, type=float,
+parser.add_argument('--threshold', default=0.6, type=float,
                     help='Only boxes with the score larger than this will be detected.')
 parser.add_argument('--use_tta', dest='use_tta', action='store_true',
                     help='Either use TTA to help detect images.')