Car Parts Segmentation App

.gitignore

@@ -0,0 +1,3 @@
+requirements.in
+__pycache__
+weight/weight.pt

app.py

@@ -0,0 +1,41 @@
+import os
+import gdown
+import gradio as gr
+from utils import Predictor
+
+PT_URL = "https://drive.google.com/uc?id=1I_LKds9obElNIZcW_DM8zyknrwRmrASj"
+PT_OUT = "weight/weights.pt"
+
+SAMPLE = [
+    [f"{root}/{file}"] \
+        for root, _, files in os.walk("sample", topdown=False) \
+        for file in files
+]
+
+if not os.path.exists(PT_OUT):
+    gdown.download(PT_URL, PT_OUT, quiet=True)
+    
+predictor = Predictor(PT_OUT)
+
+box_ui = gr.Interface(
+    fn=predictor.annotate_boxes,
+    inputs=[ gr.components.Image(type="filepath", label="Input Image") ],
+    outputs=[ gr.components.Image(type="numpy", label="Output Image") ],
+    title="Car parts segmentation",
+    examples=SAMPLE,
+    cache_examples=False,
+)
+
+mask_ui = gr.Interface(
+    fn=predictor.annotate_masks,
+    inputs=[ gr.components.Image(type="filepath", label="Input Image") ],
+    outputs=[ gr.components.Image(type="numpy", label="Output Image") ],
+    title="Car parts segmentation",
+    examples=SAMPLE,
+    cache_examples=False,
+)
+
+gr.TabbedInterface(
+    [box_ui, mask_ui],
+    tab_names=['Bouding Boxes inference', "Masks inference"]
+).queue().launch()

requirements.txt

@@ -0,0 +1,276 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    pip-compile requirements.in
+#
+aiofiles==23.2.0
+    # via gradio
+aiohttp==3.8.5
+    # via gradio
+aiosignal==1.3.1
+    # via aiohttp
+altair==5.0.1
+    # via gradio
+annotated-types==0.5.0
+    # via pydantic
+anyio==3.7.1
+    # via
+    #   httpcore
+    #   starlette
+async-timeout==4.0.2
+    # via aiohttp
+attrs==23.1.0
+    # via
+    #   aiohttp
+    #   jsonschema
+    #   referencing
+beautifulsoup4==4.12.2
+    # via gdown
+certifi==2023.7.22
+    # via
+    #   httpcore
+    #   httpx
+    #   requests
+charset-normalizer==3.2.0
+    # via
+    #   aiohttp
+    #   requests
+click==8.1.6
+    # via uvicorn
+contourpy==1.1.0
+    # via matplotlib
+cycler==0.11.0
+    # via matplotlib
+exceptiongroup==1.1.2
+    # via anyio
+fastapi==0.101.0
+    # via gradio
+ffmpy==0.3.1
+    # via gradio
+filelock==3.12.2
+    # via
+    #   gdown
+    #   huggingface-hub
+    #   torch
+fonttools==4.42.0
+    # via matplotlib
+frozenlist==1.4.0
+    # via
+    #   aiohttp
+    #   aiosignal
+fsspec==2023.6.0
+    # via
+    #   gradio-client
+    #   huggingface-hub
+gdown==4.7.1
+    # via -r requirements.in
+gradio==3.39.0
+    # via -r requirements.in
+gradio-client==0.3.0
+    # via gradio
+h11==0.14.0
+    # via
+    #   httpcore
+    #   uvicorn
+httpcore==0.17.3
+    # via httpx
+httpx==0.24.1
+    # via
+    #   gradio
+    #   gradio-client
+huggingface-hub==0.16.4
+    # via
+    #   gradio
+    #   gradio-client
+idna==3.4
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+    #   yarl
+importlib-resources==6.0.1
+    # via matplotlib
+jinja2==3.1.2
+    # via
+    #   altair
+    #   gradio
+    #   torch
+jsonschema==4.19.0
+    # via altair
+jsonschema-specifications==2023.7.1
+    # via jsonschema
+kiwisolver==1.4.4
+    # via matplotlib
+linkify-it-py==2.0.2
+    # via markdown-it-py
+markdown-it-py[linkify]==2.2.0
+    # via
+    #   gradio
+    #   mdit-py-plugins
+markupsafe==2.1.3
+    # via
+    #   gradio
+    #   jinja2
+matplotlib==3.7.2
+    # via
+    #   gradio
+    #   seaborn
+    #   ultralytics
+mdit-py-plugins==0.3.3
+    # via gradio
+mdurl==0.1.2
+    # via markdown-it-py
+mpmath==1.3.0
+    # via sympy
+multidict==6.0.4
+    # via
+    #   aiohttp
+    #   yarl
+networkx==3.1
+    # via torch
+numpy==1.25.2
+    # via
+    #   -r requirements.in
+    #   altair
+    #   contourpy
+    #   gradio
+    #   matplotlib
+    #   opencv-python
+    #   pandas
+    #   scipy
+    #   seaborn
+    #   torchvision
+    #   ultralytics
+opencv-python==4.8.0.74
+    # via ultralytics
+orjson==3.9.4
+    # via gradio
+packaging==23.1
+    # via
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   matplotlib
+pandas==2.0.3
+    # via
+    #   -r requirements.in
+    #   altair
+    #   gradio
+    #   seaborn
+    #   ultralytics
+pillow==10.0.0
+    # via
+    #   gradio
+    #   matplotlib
+    #   torchvision
+    #   ultralytics
+psutil==5.9.5
+    # via ultralytics
+py-cpuinfo==9.0.0
+    # via ultralytics
+pydantic==2.1.1
+    # via
+    #   fastapi
+    #   gradio
+pydantic-core==2.4.0
+    # via pydantic
+pydub==0.25.1
+    # via gradio
+pyparsing==3.0.9
+    # via matplotlib
+pysocks==1.7.1
+    # via requests
+python-dateutil==2.8.2
+    # via
+    #   matplotlib
+    #   pandas
+python-multipart==0.0.6
+    # via gradio
+pytz==2023.3
+    # via pandas
+pyyaml==6.0.1
+    # via
+    #   gradio
+    #   huggingface-hub
+    #   ultralytics
+referencing==0.30.2
+    # via
+    #   jsonschema
+    #   jsonschema-specifications
+requests[socks]==2.31.0
+    # via
+    #   gdown
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   torchvision
+    #   ultralytics
+rpds-py==0.9.2
+    # via
+    #   jsonschema
+    #   referencing
+scipy==1.11.1
+    # via ultralytics
+seaborn==0.12.2
+    # via ultralytics
+semantic-version==2.10.0
+    # via gradio
+six==1.16.0
+    # via
+    #   gdown
+    #   python-dateutil
+sniffio==1.3.0
+    # via
+    #   anyio
+    #   httpcore
+    #   httpx
+soupsieve==2.4.1
+    # via beautifulsoup4
+starlette==0.27.0
+    # via fastapi
+sympy==1.12
+    # via torch
+toolz==0.12.0
+    # via altair
+torch==2.0.1
+    # via
+    #   torchvision
+    #   ultralytics
+torchvision==0.15.2
+    # via ultralytics
+tqdm==4.66.0
+    # via
+    #   gdown
+    #   huggingface-hub
+    #   ultralytics
+typing-extensions==4.7.1
+    # via
+    #   altair
+    #   fastapi
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   pydantic
+    #   pydantic-core
+    #   starlette
+    #   torch
+    #   uvicorn
+tzdata==2023.3
+    # via pandas
+uc-micro-py==1.0.2
+    # via linkify-it-py
+ultralytics==8.0.150
+    # via -r requirements.in
+urllib3==2.0.4
+    # via requests
+uvicorn==0.23.2
+    # via gradio
+websockets==11.0.3
+    # via
+    #   gradio
+    #   gradio-client
+yarl==1.9.2
+    # via aiohttp
+zipp==3.16.2
+    # via importlib-resources

utils/__init__.py

@@ -0,0 +1,89 @@
+import cv2
+import numpy as np 
+import matplotlib.pyplot as plt
+from matplotlib import colors
+from ultralytics import YOLO
+from ultralytics.utils.ops import scale_image
+
+class Predictor:
+
+  def __init__(self, model_weight):
+    self.model = YOLO(model_weight)
+    self.category_map = self.model.names
+    self.NCLS = len(self.category_map)
+    
+    cmap = plt.cm.rainbow
+    cmaplist = [cmap(i) for i in range(cmap.N)]
+
+    self.cmap = cmap.from_list(f'my cmap', cmaplist, cmap.N)
+
+    bounds = np.linspace(0, self.NCLS, self.NCLS + 1)
+    norm = colors.BoundaryNorm(bounds, self.cmap.N)
+
+    category_cmap = { k: cmap(norm(int(k))) for k in self.category_map }
+    self.category_cmap = { k: (v[2] * 255, v[1] * 255, v[0]* 255) \
+                          for k, v in category_cmap.items() }
+
+  def predict(self, image_path):
+    image = cv2.imread(image_path)
+    outputs = self.model.predict(source=image_path)
+    results = outputs[0].cpu().numpy()
+
+    boxes = results.boxes.xyxy
+    confs = results.boxes.conf
+    cls   = results.boxes.cls
+    # probs = results.boxes.probs
+    masks = results.masks.data
+
+    return image, cls, confs, boxes, masks, results
+
+  def annotate_boxes(self, image_path):
+    image, cls, confs, boxes, _, results = self.predict(image_path)
+
+    for i, (box, cl, conf) in enumerate(zip(boxes, cls, confs)):
+      label = results.names[cl]
+      color = self.category_cmap[cl]
+      text = label + f" {conf:.2f}"
+      x1, y1, x2, y2 = ( int(p) for p in box )
+
+      cv2.rectangle(image, (x1, y1), (x2, y2),
+                    color=color,
+                    thickness=2,
+                    lineType=cv2.LINE_AA
+                    )
+      (w, h), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_DUPLEX, 0.3, 1)
+      cv2.rectangle(image, (x1, y1 - 2*h), (x1 + w, y1), color, -1)
+      cv2.putText(image, text, (x1, y1 - 5), 
+                    cv2.FONT_HERSHEY_DUPLEX, 0.3, (255, 255, 255), 1)
+
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+  def annotate_masks(self, image_path):
+    image, cls, confs, _, masks, results = self.predict(image_path)
+    ori_shape = image.shape[:2]
+
+    for i, (mask, cl, conf) in enumerate(zip(masks, cls, confs)):
+      mask = mask.astype("uint8")
+      label = results.names[cl]
+      color = self.category_cmap[cl]
+      text = label + f" {conf:.2f}"
+
+      _mask = np.where(mask[..., None], color, (0, 0, 0))
+      _mask = scale_image(_mask, ori_shape).astype("uint8")
+      image = cv2.addWeighted(image, 1, _mask, 0.5, 0)
+
+      contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+      boundary = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGBA).astype("float")
+      cv2.drawContours(boundary, contours, -1, color, 2)
+      boundary = scale_image(boundary, ori_shape)[:, :, :-1].astype("uint8")
+      image = cv2.addWeighted(image, 1, boundary, 1, 0)
+
+      cy, cx = np.round(np.argwhere(_mask != [0, 0, 0]).mean(axis=0)[0:2]).astype(int)
+      (w, h), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_DUPLEX, 0.5, 1)
+
+      cv2.putText(image, text, (cx - int(0.5 * w), cy),
+                cv2.FONT_HERSHEY_DUPLEX, 0.5, (0, 0, 0), 2)
+      cv2.putText(image, text, (cx - int(0.5 * w), cy),
+                cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255), 1)
+      
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)