Add crop image based on class detected

data/data.yaml

@@ -0,0 +1,5 @@
+train: ../train/images
+val: ../valid/images
+
+nc: 2
+names: ['license-plate', 'vehicle']

detect.py

@@ -24,6 +24,14 @@
                                          yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
 """
 
+from utils.openalpr import Alpr
+from utils.torch_utils import select_device, time_sync
+from utils.plots import Annotator, colors, save_one_box
+from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
+                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
+from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
+from models.common import DetectMultiBackend
+from utils.image_crop import crop
 import argparse
 import os
 import sys
@@ -38,13 +46,6 @@
     sys.path.append(str(ROOT))  # add ROOT to PATH
 ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
 
-from models.common import DetectMultiBackend
-from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
-from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
-                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
-from utils.plots import Annotator, colors, save_one_box
-from utils.torch_utils import select_device, time_sync
-
 
 @torch.no_grad()
 def run(
@@ -151,24 +152,30 @@ def run(
                 det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
 
                 # Print results
-                for c in det[:, -1].unique():
-                    n = (det[:, -1] == c).sum()  # detections per class
-                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+                # for c in det[:, -1].unique():
+                #     n = (det[:, -1] == c).sum()  # detections per class
+                #     s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
 
                 # Write results
                 for *xyxy, conf, cls in reversed(det):
-                    if save_txt:  # Write to file
-                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
-                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
-                        with open(f'{txt_path}.txt', 'a') as f:
-                            f.write(('%g ' * len(line)).rstrip() % line + '\n')
-
-                    if save_img or save_crop or view_img:  # Add bbox to image
-                        c = int(cls)  # integer class
-                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
-                        annotator.box_label(xyxy, label, color=colors(c, True))
-                    if save_crop:
-                        save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
+                    # if save_txt:  # Write to file
+                    #     xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
+                    #     line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
+                    #     with open(f'{txt_path}.txt', 'a') as f:
+                    #         f.write(('%g ' * len(line)).rstrip() % line + '\n')
+
+                    # if save_img or save_crop or view_img:  # Add bbox to image
+                    #     c = int(cls)  # integer class
+                    #     label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
+                    #     annotator.box_label(xyxy, label, color=colors(c, True))
+                    # if save_crop:
+                    #     save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
+                    c = int(cls)
+                    if c == 0:
+                        label = None
+                        imCrop = crop(xyxy, imc)
+                        cv2.imshow("cropped", imCrop)
+                        cv2.waitKey(1)
 
             # Stream results
             im0 = annotator.result()

utils/image_crop.py

@@ -0,0 +1,28 @@
+import cv2
+import torch
+import numpy as np
+
+from utils.general import (xywh2xyxy, xyxy2xywh)
+
+
+def crop(xyxy, im, square=False, gain=1.02, pad=10, BGR=False):
+    xyxy = torch.tensor(xyxy).view(-1, 4)
+
+    b = xyxy2xywh(xyxy)
+
+    if square:
+        b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1)
+
+    b[:, 2:] = b[:, 2:] * gain + pad
+
+    xyxy = xywh2xyxy(b).long()
+
+    crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
+
+    gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY)
+
+    medBlur = cv2.medianBlur(gray, 3)
+
+    gaussBlur = cv2.GaussianBlur(medBlur, (5, 5), 0)
+
+    return gaussBlur