Skip to content

Commit

Permalink
update cutout script
Browse files Browse the repository at this point in the history
  • Loading branch information
@capjamesg
capjamesg committed May 24, 2023
1 parent fcb878b commit 3dd2424
Showing 1 changed file with 109 additions and 104 deletions.
213 changes: 109 additions & 104 deletions scripts/cutout.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,17 +31,16 @@
parser.add_argument(
"--generate_pdf", type=bool, required=False, default=False, help="Generate PDF"
)
parser.add_argument(
"--interactive", type=bool, required=False, default=False, help="Show false positives"
)
parser.add_argument(
"--fp_threshold",
type=float,
required=False,
default=0.7,
help="False positive threshold",
)
# iou threshold
parser.add_argument(
"--iou_threshold", type=float, required=False, default=0.5, help="IOU threshold"
)

args = parser.parse_args()

Expand All @@ -59,6 +58,14 @@
dataset="train"
).download_dataset()


mask_vectors = {}
clip_vectors_by_class = {}
bbox_size = []
class_distribution = {}
whole_image_vectors = []
vectors_by_image = {}

images = data.keys()

device = "cuda" if torch.cuda.is_available() else "cpu"
Expand All @@ -76,105 +83,71 @@ def run_clip_inference(mask: str) -> torch.Tensor:
return image_features


mask_vectors = {}
clip_vectors_by_class = {}
bbox_size = []
class_distribution = {}
whole_image_vectors = []
vectors_by_image = {}
def save_report():
print("Computing images with potential false positives. This may take some time, depending on with how many images you are working.")

for i in data.keys():
img = cv2.imread(i)
print("Computing ground truth vectors for", i)
potential_false_positives = []

for j in data[i]["ground_truth"]:
x1 = int(j[0])
y1 = int(j[1])
x2 = int(j[2])
y2 = int(j[3])
class_name = class_names[int(j[4])]
mask = img[y1:y2, x1:x2]
for i in data.keys():
img = cv2.imread(i)

if class_name not in clip_vectors_by_class:
clip_vectors_by_class[class_name] = []
print("Evaluating", i)

vector = run_clip_inference(mask)
for j in data[i]["ground_truth"]:
x1 = int(j[0])
y1 = int(j[1])
x2 = int(j[2])
y2 = int(j[3])
class_name = class_names[int(j[4])]

class_distribution[class_name] = class_distribution.get(class_name, 0) + 1
if (x1, y1, x2, y2) not in mask_vectors[class_name]:
continue

if mask_vectors.get(class_name) is None:
mask_vectors[class_name] = {}
img = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)

mask_vectors[class_name][(x1, y1, x2, y2)] = vector
clip_vector = mask_vectors[class_name][(x1, y1, x2, y2)]

clip_vectors_by_class[class_name].append(vector)
bbox_size.append(mask.shape[0] * mask.shape[1])
whole_image_clip = run_clip_inference(img)
# show mask
mask = img[y1:y2, x1:x2]

whole_image_vectors.append(whole_image_clip)
vectors_by_image[i] = whole_image_clip
# calculate distance between clip vector and class vector
distance = torch.dist(clip_vector, clip_vectors_by_class[class_name])

for i in clip_vectors_by_class.keys():
clip_vectors_by_class[i] = torch.mean(torch.stack(clip_vectors_by_class[i]), dim=0)
if distance > 0.7:
print(f"{i} has a false positive at {x1, y1, x2, y2}")
image = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv2.imshow("image", image)

whole_image_vectors = torch.mean(torch.stack(whole_image_vectors), dim=0)
potential_false_positives.append(i)

avg_frame_size = sum(bbox_size) / len(bbox_size)
mega_pixels = round(avg_frame_size / 1000000, 2)
if args.interactive:
cv2.imshow("mask", mask)

class_distribution_as_str = """
| Class Name | Count |
| ---------- | ----- |
{}""".format(
"\n".join(
[
"| {} | {} |".format(class_name, count)
for class_name, count in class_distribution.items()
]
)
)
if cv2.waitKey(0) == ord("q"):
args.interactive = False

images_with_potential_false_positives = f"- {os.linesep.join([i for i in data.keys() if torch.dist(run_clip_inference(cv2.imread(i)), whole_image_vectors) > args.fp_threshold])}"
cv2.destroyAllWindows()

potential_false_positives = "\n\n".join(potential_false_positives)

def save_report():
report = f"""
# Bounding Box Size Report
## Stats
- Dataset Used: {ROBOFLOW_WORKSPACE_URL}/{ROBOFLOW_PROJECT_URL}/{ROBOFLOW_MODEL_VERSION}
- Dataset Used: {ROBOFLOW_PROJECT_URL}/1
- Number of Images: {len(images)}
- Number of Classes: {len(class_names)}
## Average Bounding Box Size
The average bounding box size is {mega_pixels} megapixels.
## Class Distribution
{class_distribution_as_str}
## Images with Potential False Positives
{images_with_potential_false_positives}
{potential_false_positives}
"""

with open("fps.csv", "w") as f:
writer = csv.writer(f)

writer.writerow(["image", "fps"])

for i in data.keys():
writer.writerow(
[
i,
torch.dist(
run_clip_inference(cv2.imread(i)), whole_image_vectors
).item(),
]
)

# save to file

with open("report.md", "w") as f:
f.write(report)

Expand All @@ -186,38 +159,12 @@ def save_report():
extra_args=["-V", "geometry:margin=1in", "-V", "mainfont=Times New Roman"],
)

for i in data.keys():
img = cv2.imread(i)

print("Evaluating", i)

for j in data[i]["ground_truth"]:
x1 = int(j[0])
y1 = int(j[1])
x2 = int(j[2])
y2 = int(j[3])
class_name = class_names[int(j[4])]
img = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)

clip_vector = mask_vectors[class_name][(x1, y1, x2, y2)]

# show mask
mask = img[y1:y2, x1:x2]

# calculate distance between clip vector and class vector
distance = torch.dist(clip_vector, clip_vectors_by_class[class_name])

print(distance)
if distance > 0.7:
print("false positive")
image = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv2.imshow("image", image)

if cv2.waitKey(0) & 0xFF == ord("q"):
break

# find duplicate clip vectors
def find_duplicate_images():
# show progress bar

fps = 0

images = os.listdir(EVAL_DATA_PATH + "train/images")

for i in images:
Expand All @@ -234,7 +181,65 @@ def find_duplicate_images():
if distance < 0.1:
print(f"{i} and {j} are duplicates")

for i in data.keys():
img = cv2.imread(i)
print("Computing ground truth vectors for", i)

for j in data[i]["ground_truth"]:
x1 = int(j[0])
y1 = int(j[1])
x2 = int(j[2])
y2 = int(j[3])
class_name = class_names[int(j[4])]

if x1 == x2 or y1 == y2:
print(f"{i} has a 0x0 bounding box")
continue

mask = img[y1:y2, x1:x2]

if class_name not in clip_vectors_by_class:
clip_vectors_by_class[class_name] = []

vector = run_clip_inference(mask)

class_distribution[class_name] = class_distribution.get(class_name, 0) + 1

if mask_vectors.get(class_name) is None:
mask_vectors[class_name] = {}

mask_vectors[class_name][(x1, y1, x2, y2)] = vector

clip_vectors_by_class[class_name].append(vector)
bbox_size.append(mask.shape[0] * mask.shape[1])
whole_image_clip = run_clip_inference(img)

whole_image_vectors.append(whole_image_clip)
vectors_by_image[i] = whole_image_clip

print("Computing mean class vectors")

for i in clip_vectors_by_class.keys():
clip_vectors_by_class[i] = torch.mean(torch.stack(clip_vectors_by_class[i]), dim=0)

whole_image_vectors = torch.mean(torch.stack(whole_image_vectors), dim=0)

avg_frame_size = sum(bbox_size) / len(bbox_size)
mega_pixels = round(avg_frame_size / 1000000, 2)

class_distribution_as_str = """
| Class Name | Count |
| ---------- | ----- |
{}""".format(
"\n".join(
[
"| {} | {} |".format(class_name, count)
for class_name, count in class_distribution.items()
]
)
)

if __name__ == "__main__":
find_duplicate_images()
# find_duplicate_images()
print("Saving report")
save_report()

0 comments on commit 3dd2424

Please sign in to comment.