Update Damage Severity Detection

app.py

@@ -1,10 +1,11 @@
 import gradio as gr
 from utils import Predictor, setup
 
-CAR_PART_SEG_OUT, CAR_DAM_DET_OUT, SAMPLE = setup()
+CAR_PART_SEG_OUT, CAR_DAM_DET_OUT, CAR_SEV_DET_OUT, SAMPLE = setup()
 
 car_part_predictor = Predictor(CAR_PART_SEG_OUT)
 car_dam_detector = Predictor(CAR_DAM_DET_OUT)
+car_sev_detector = Predictor(CAR_SEV_DET_OUT)
 
 car_part_ui = gr.Interface(
     fn=car_part_predictor.transform,
@@ -32,7 +33,24 @@ car_dam_ui = gr.Interface(
     cache_examples=False,
 )
 
+car_sev_ui = gr.Interface(
+    fn=car_sev_detector.transform,
+    inputs=[ 
+        gr.components.Image(type="filepath", label="Input Image"),
+        gr.components.Checkbox(value=True, label="Annotate Boxes"),
+        gr.components.Checkbox(value=False, label="Annotate Masks"),
+        ],
+    outputs=gr.components.Image(type="numpy", label="Output Image"),
+    title="Car Damage Severity Detector",
+    # examples=SAMPLE["car-dam-det"],
+    cache_examples=False,
+)
+
 gr.TabbedInterface(
-    [car_part_ui, car_dam_ui],
-    tab_names=['car part segmentation', 'car damage detection']
+    [car_part_ui, car_dam_ui, car_sev_ui],
+    tab_names=[
+        'car part segmentation', 
+        'car damage detection', 
+        'car damage severity detection'
+        ]
 ).queue().launch()

utils/__init__.py

@@ -10,9 +10,11 @@ from ultralytics.utils.ops import scale_image
 def setup():
   CAR_PART_SEG_URL = "https://drive.google.com/uc?id=1I_LKds9obElNIZcW_DM8zyknrwRmrASj"
   CAR_DAM_DET_URL = "https://drive.google.com/uc?id=1AXDyFoEuNqXSaDNUHBp8H9AVjICvsUpz"
+  CAR_SEV_DET_URL = "https://drive.google.com/uc?id=1An7QGjbL-UEu7LOT7Xh59itBE854vy4U"
 
   CAR_PART_SEG_OUT = "weight/yolov8-car-part-seg.pt"
   CAR_DAM_DET_OUT = "weight/yolov8-car-damage-detection.pt"
+  CAR_SEV_DET_OUT = "weight/yolov8-car-damage-serverity-detection.pt"
 
   SAMPLE = {
       "car-parts-seg" : [ [f"{root}/{file}"] \
@@ -30,8 +32,12 @@ def setup():
   if not os.path.exists(CAR_DAM_DET_OUT):
       os.makedirs("weight", exist_ok=True)
       gdown.download(CAR_DAM_DET_URL, CAR_DAM_DET_OUT, quiet=True)
+      
+  if not os.path.exists(CAR_SEV_DET_URL):
+      os.makedirs("weight", exist_ok=True)
+      gdown.download(CAR_SEV_DET_URL, CAR_SEV_DET_OUT, quiet=True)
   
-  return CAR_PART_SEG_OUT, CAR_DAM_DET_OUT, SAMPLE
+  return CAR_PART_SEG_OUT, CAR_DAM_DET_OUT, CAR_SEV_DET_OUT, SAMPLE
 
 class Predictor: