Add CoreML iOS updates (#121)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>

ultralytics/yolo/utils/checks.py

@@ -22,16 +22,26 @@ def is_ascii(s=''):
     return len(s.encode().decode('ascii', 'ignore')) == len(s)
 
 
-def check_imgsz(imgsz, s=32, floor=0):
+def check_imgsz(imgsz, stride=32, min_dim=1, floor=0):
     # Verify image size is a multiple of stride s in each dimension
-    if isinstance(imgsz, int):  # integer i.e. img_size=640
-        new_size = max(make_divisible(imgsz, int(s)), floor)
-    else:  # list i.e. img_size=[640, 480]
+
+    stride = int(stride.max() if isinstance(stride, torch.Tensor) else stride)
+    if isinstance(imgsz, int):  # integer i.e. imgsz=640
+        sz = max(make_divisible(imgsz, stride), floor)
+    else:  # list i.e. imgsz=[640, 480]
         imgsz = list(imgsz)  # convert to list if tuple
-        new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz]
-    if new_size != imgsz:
-        LOGGER.warning(f'WARNING ⚠️ --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}')
-    return new_size
+        sz = [max(make_divisible(x, stride), floor) for x in imgsz]
+    if sz != imgsz:
+        LOGGER.warning(f'WARNING ⚠️ --img-size {imgsz} must be multiple of max stride {stride}, updating to {sz}')
+
+    # Check dims
+    if min_dim == 2:
+        if isinstance(imgsz, int):
+            sz = [sz, sz]
+        elif len(sz) == 1:
+            sz = [sz[0], sz[0]]
+
+    return sz
 
 
 def check_version(current="0.0.0", minimum="0.0.0", name="version ", pinned=False, hard=False, verbose=False):

ultralytics/yolo/utils/torch_utils.py

@@ -185,18 +185,6 @@ def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
     return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
 
 
-def check_imgsz(imgsz, s=32, floor=0):
-    # Verify image size is a multiple of stride s in each dimension
-    if isinstance(imgsz, int):  # integer i.e. imgsz=640
-        new_size = max(make_divisible(imgsz, int(s)), floor)
-    else:  # list i.e. imgsz=[640, 480]
-        imgsz = list(imgsz)  # convert to list if tuple
-        new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz]
-    if new_size != imgsz:
-        LOGGER.warning(f'WARNING ⚠️ --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}')
-    return new_size
-
-
 def make_divisible(x, divisor):
     # Returns nearest x divisible by divisor
     if isinstance(divisor, torch.Tensor):
@@ -293,3 +281,18 @@ def strip_optimizer(f='best.pt', s=''):  # from utils.general import *; strip_op
     torch.save(x, s or f)
     mb = os.path.getsize(s or f) / 1E6  # filesize
     LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
+
+
+def guess_task_from_head(head):
+    task = None
+    if head.lower() in ["classify", "classifier", "cls", "fc"]:
+        task = "classify"
+    if head.lower() in ["detect"]:
+        task = "detect"
+    if head.lower() in ["segment"]:
+        task = "segment"
+
+    if not task:
+        raise SyntaxError("task or model not recognized! Please refer the docs at : ")  # TODO: add docs links
+
+    return task