target per shape.py 2.26 KB
import cv2
import os
import numpy as np
import shutil 
import time

font = cv2.FONT_HERSHEY_COMPLEX

load_path = "./total_frame_m1"
#load_path = "./sample"
color = (255,255,255)

for path, dirs, files in os.walk(load_path):
    for filename in files:
        img = cv2.imread(os.path.join(path, filename), cv2.IMREAD_GRAYSCALE)
        # (img, threshold_value, value, flag): if pixel val >= threshold_value: pixel_val = value
        _, threshold = cv2.threshold(img, 1, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(threshold, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        n_ellipse, n_circle = 0, 0
        if (len(contours) > 0):
            area = cv2.contourArea(contours[0])
        largest = ""

        for cnt in contours:
            approx = cv2.approxPolyDP(cnt, 0.01*cv2.arcLength(cnt, True), True)
            cv2.drawContours(img, [approx], 0, (0), 5)
            x = approx.ravel()[0]
            y = approx.ravel()[1]

            #print("approx", len(approx))
            if len(approx) < 12:
                #cv2.putText(img, "Ellipse", (x, y), font, 1, color)
                n_ellipse = n_ellipse + 1
                if (area < cv2.contourArea(cnt)):
                    area = cv2.contourArea(cnt)
                    largest = "e"
            else:
                #cv2.putText(img, "Circle", (x, y), font, 1, color)
                n_circle = n_circle + 1
                if (area < cv2.contourArea(cnt)):
                    area = cv2.contourArea(cnt)
                    largest = "c"

        print(n_ellipse, n_circle, filename)
        #assert n_ellipse !=0 or n_circle !=0
        
        if (n_ellipse > n_circle) or (n_ellipse==n_circle and largest == "e"):
            ellipase_path = os.path.join("./ellipse/", filename)
            shutil.copy(os.path.join(path, filename), ellipase_path)

        elif (n_ellipse < n_circle) or (n_ellipse==n_circle and largest == "c"):
            circle_path = os.path.join("./circle/", filename)
            shutil.copy(os.path.join(path, filename), circle_path)

        elif n_circle == 0 and n_ellipse == 0:
            shutil.copy(os.path.join(path, filename), os.path.join("./zero/", filename))
        

# cv2.imshow("shapes", img)
# cv2.imshow("Threshold", threshold)
# cv2.waitKey(0)
# cv2.destroyAllWindows()