Samll update
This commit is contained in:
@@ -18,6 +18,8 @@ import argparse
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import random
|
import random
|
||||||
|
|
||||||
|
from src.utils.image import Image
|
||||||
|
|
||||||
|
|
||||||
def parse_label_line(line):
|
def parse_label_line(line):
|
||||||
parts = line.strip().split()
|
parts = line.strip().split()
|
||||||
@@ -62,7 +64,11 @@ def poly_to_pts(coords, img_w, img_h):
|
|||||||
|
|
||||||
def random_color_for_class(cls):
|
def random_color_for_class(cls):
|
||||||
random.seed(cls) # deterministic per class
|
random.seed(cls) # deterministic per class
|
||||||
return tuple(int(x) for x in np.array([random.randint(0, 255) for _ in range(3)]))
|
return (
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
255,
|
||||||
|
) # tuple(int(x) for x in np.array([random.randint(0, 255) for _ in range(3)]))
|
||||||
|
|
||||||
|
|
||||||
def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
|
def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
|
||||||
@@ -77,13 +83,13 @@ def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
|
|||||||
color = random_color_for_class(cls)
|
color = random_color_for_class(cls)
|
||||||
|
|
||||||
x1, y1, x2, y2 = yolo_bbox_to_xyxy(coords[:4], w, h)
|
x1, y1, x2, y2 = yolo_bbox_to_xyxy(coords[:4], w, h)
|
||||||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)
|
||||||
|
|
||||||
pts = poly_to_pts(coords[4:], w, h)
|
pts = poly_to_pts(coords[4:], w, h)
|
||||||
# fill on overlay
|
# fill on overlay
|
||||||
cv2.fillPoly(overlay, [pts], color)
|
cv2.fillPoly(overlay, [pts], color)
|
||||||
# outline on base image
|
# outline on base image
|
||||||
cv2.polylines(img, [pts], isClosed=True, color=color, thickness=2)
|
cv2.polylines(img, [pts], isClosed=True, color=color, thickness=1)
|
||||||
# put class text at first point
|
# put class text at first point
|
||||||
x, y = int(pts[0, 0]), int(pts[0, 1]) - 6
|
x, y = int(pts[0, 0]), int(pts[0, 1]) - 6
|
||||||
cv2.putText(
|
cv2.putText(
|
||||||
@@ -158,7 +164,9 @@ def main():
|
|||||||
print("Label file not found:", lbl_path)
|
print("Label file not found:", lbl_path)
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
|
# img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
|
||||||
|
img = (Image(img_path).get_qt_rgb() * 255).astype(np.uint8)
|
||||||
|
|
||||||
if img is None:
|
if img is None:
|
||||||
print("Could not load image:", img_path)
|
print("Could not load image:", img_path)
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
@@ -173,6 +181,7 @@ def main():
|
|||||||
|
|
||||||
# Convert BGR -> RGB for matplotlib display
|
# Convert BGR -> RGB for matplotlib display
|
||||||
out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
|
out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
|
||||||
|
# out_rgb = Image()
|
||||||
plt.figure(figsize=(10, 10 * out.shape[0] / out.shape[1]))
|
plt.figure(figsize=(10, 10 * out.shape[0] / out.shape[1]))
|
||||||
plt.imshow(out_rgb)
|
plt.imshow(out_rgb)
|
||||||
plt.axis("off")
|
plt.axis("off")
|
||||||
|
|||||||
Reference in New Issue
Block a user